Размер от in_channels
до self.fc1
зависит от размера входного изображения, а не от размера ядра.
В вашем случае self.fc1 = nn.Linear(16 * 5 * 5, 120)
должно быть nn.Linear(16 * image_size * image_size)
где, image_size
: размер изображения в последнем слое свертки.
Пример кода:
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
class Net(nn.Module):
def __init__(self, classes):
super().__init__()
self.conv1 = nn.Conv2d(3, 6, 5, padding=2)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5, padding=2)
self.fc1 = nn.Linear(16 * 25 * 25, 120)
self.fc2 = nn.Linear(120, 50)
self.fc3 = nn.Linear(50, classes)
def forward(self, x):
print('one', x.shape)
x = self.pool(F.relu(self.conv1(x)))
print('two', x.shape)
x = self.pool(F.relu(self.conv2(x)))
print('three', x.shape)
x = x.view(-1, np.product(x.shape[1:]))
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
x = torch.rand((32, 3, 100, 100))
net = Net(2)
out= net(x)
print('out', out.shape)
one torch.Size([32, 3, 100, 100])
two torch.Size([32, 6, 50, 50])
three torch.Size([32, 16, 25, 25])
out torch.Size([32, 2])