|
|
|
@@ -31,20 +31,25 @@ class LeNet5(nn.Cell): |
|
|
|
>>> LeNet(num_class=10) |
|
|
|
|
|
|
|
""" |
|
|
|
def __init__(self, num_class=10, num_channel=1): |
|
|
|
def __init__(self, num_class=10, num_channel=1, include_top=True): |
|
|
|
super(LeNet5, self).__init__() |
|
|
|
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, 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() |
|
|
|
self.include_top = include_top |
|
|
|
if self.include_top: |
|
|
|
self.flatten = nn.Flatten() |
|
|
|
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)) |
|
|
|
|
|
|
|
|
|
|
|
def construct(self, x): |
|
|
|
x = self.max_pool2d(self.relu(self.conv1(x))) |
|
|
|
x = self.max_pool2d(self.relu(self.conv2(x))) |
|
|
|
if not self.include_top: |
|
|
|
return x |
|
|
|
x = self.flatten(x) |
|
|
|
x = self.relu(self.fc1(x)) |
|
|
|
x = self.relu(self.fc2(x)) |
|
|
|
|