一般这种问题是由于输入图像的大小不匹配导致,
class Network(nn.Module):
def __init__(self, height_size, width_size, num_classes):
super(Network, self).__init__()
self.height_size = height_size
self.width_size = width_size
self.conv2d_1 = nn.Conv2d(in_channels=3, out_channels=8, kernel_size=(3, 3), padding=2)
self.pool = nn.MaxPool2d(2, 2)
self.conv2d_2 = nn.Conv2d(in_channels=8, out_channels=16, kernel_size=(3, 3), padding=2)
#self.fc_1 = nn.Linear(16 * self.height_size // 4*self.width_size // 4, 1024)
self.fc_1 = nn.Linear(16 * 9 * 9, 1024)
self.fc_2 = nn.Linear(1024, num_classes)
def forward(self, input):
x = self.conv2d_1(input)
x = F.relu(x)
x = self.pool(x)
x = self.conv2d_2(x)
x = F.relu(x)
x = self.pool(x)
#x = x.view(-1, 16*self.height_size//4*self.width_size//4)
#print(x.shape)
#torch.Tensor(64, 16, 9, 9)
x = x.view(-1, 16 * 9 * 9)
x = self.fc_1(x)
x = F.relu(x)
x = F.dropout(x, training=self.training)
x = self.fc_2(x)
output = F.log_softmax(x,dim=1)
return output
作者这里是由于修改了kernal_size导致卷积之后的图像大小与修改之前的图像大小不一致报错。
可以通过print一下view之前的输入形状并根据结果来修改线性层的输入大小,记得线性层初始化时的输入大小也要修改。