p83571920
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

modify lenet normal

This commit is contained in:
wukesong 2020-09-03 21:33:29 +08:00
parent 4283c88e0f
commit fb568c7dba
1 changed files with 7 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
model_zoo/official/cv/lenet/src/lenet.py
View File

@ -14,6 +14,7 @@
# ============================================================================
"""LeNet."""
import mindspore.nn as nn
from mindspore.common.initializer import Normal
class LeNet5(nn.Cell):
@ -22,7 +23,7 @@ class LeNet5(nn.Cell):
Args:
num_class (int): Num classes. Default: 10.
channel (int): Num classes. Default: 1.
num_channel (int): Num channels. Default: 1.
Returns:
Tensor, output tensor
@ -30,14 +31,13 @@ class LeNet5(nn.Cell):
>>> LeNet(num_class=10)
"""
def __init__(self, num_class=10, channel=1):
def __init__(self, num_class=10, num_channel=1):
super(LeNet5, self).__init__()
self.num_class = num_class
self.conv1 = nn.Conv2d(channel, 6, 5, pad_mode='valid')
self.conv1 = nn.Conv2d(num_channel, 6, 5, pad_mode='valid')
self.conv2 = nn.Conv2d(6, 16, 5, pad_mode='valid')
self.fc1 = nn.Dense(16 * 5 * 5, 120)
self.fc2 = nn.Dense(120, 84)
self.fc3 = nn.Dense(84, self.num_class)
self.fc1 = nn.Dense(16 * 5 * 5, 120, weight_init=Normal(0.02))
self.fc2 = nn.Dense(120, 84, weight_init=Normal(0.02))
self.fc3 = nn.Dense(84, num_class, weight_init=Normal(0.02))
self.relu = nn.ReLU()
self.max_pool2d = nn.MaxPool2d(kernel_size=2, stride=2)
self.flatten = nn.Flatten()