1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
| class CNN(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3,6,3)
self.pool = nn.MaxPool2d(2)
self.flatten = nn.Flatten()
# After convolution and pooling, each sample has shape [6, 15, 15].
self.fc = nn.Linear(6*15*15, 10)
def forward(self, x):
# Assume the input is [B,3,32,32]
x = self.conv1(x) # [B,6,30,30] ✅ you got it right
x = self.pool(x) # [B,6,15,15] ✅ you got it right
x = self.flatten(x) # [B,6*15*15] ✅
x = self.fc(x) # [B,10]
return x
# Training
model = CNN()
criterion = nn.CrossEntropyLoss() # dedicated to classification
for inputs, labels in dataloader:
outputs = model(inputs)
loss = criterion(outputs, labels) # compute the loss
loss.backward()
|