| 173 | # |
| 174 | |
| 175 | class LeNet(nn.Module): |
| 176 | |
| 177 | def __init__(self): |
| 178 | super(LeNet, self).__init__() |
| 179 | # 1 input image channel (black & white), 6 output channels, 5x5 square convolution |
| 180 | # kernel |
| 181 | self.conv1 = nn.Conv2d(1, 6, 5) |
| 182 | self.conv2 = nn.Conv2d(6, 16, 5) |
| 183 | # an affine operation: y = Wx + b |
| 184 | self.fc1 = nn.Linear(16 * 5 * 5, 120) # 5*5 from image dimension |
| 185 | self.fc2 = nn.Linear(120, 84) |
| 186 | self.fc3 = nn.Linear(84, 10) |
| 187 | |
| 188 | def forward(self, x): |
| 189 | # Max pooling over a (2, 2) window |
| 190 | x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2)) |
| 191 | # If the size is a square you can only specify a single number |
| 192 | x = F.max_pool2d(F.relu(self.conv2(x)), 2) |
| 193 | x = x.view(-1, self.num_flat_features(x)) |
| 194 | x = F.relu(self.fc1(x)) |
| 195 | x = F.relu(self.fc2(x)) |
| 196 | x = self.fc3(x) |
| 197 | return x |
| 198 | |
| 199 | def num_flat_features(self, x): |
| 200 | size = x.size()[1:] # all dimensions except the batch dimension |
| 201 | num_features = 1 |
| 202 | for s in size: |
| 203 | num_features *= s |
| 204 | return num_features |
| 205 | |
| 206 | |
| 207 | ############################################################################ |