如何解决自定义CNN提供错误的输出形状
我需要一些帮助。我正在尝试制作一个自定义的CNN,该CNN应该接受一个通道图像并进行二进制分类。这是模型:
class custom_small_CNN(nn.Module):
def __init__(self,input_channels=1,output_features=1):
super(custom_small_CNN,self).__init__()
self.input_channels = input_channels
self.output_features = output_features
self.conv1 = nn.Conv2d(self.input_channels,8,kernel_size=(7,7),stride=(2,2),padding=(6,6),dilation=(2,2))
self.conv2 = nn.Conv2d(8,16,kernel_size=(3,3),padding=(1,1),dilation=(1,1))
self.pool = nn.MaxPool2d(kernel_size=(2,2))
self.fc1 = nn.Linear(in_features=1024,out_features=self.output_features,bias=True)
self.dropout = nn.Dropout(p=0.5)
self.softmax = nn.Softmax(dim=1)
self.net_name = 'Custom_Small_CNN'
self.net = nn.Sequential(self.conv1,self.pool,self.conv2,self.fc1)
def forward(self,x):
x = self.conv1(x)
x = self.pool(x)
#x = self.dropout(x)
x = self.conv2(x)
x = self.pool(x)
x = x.view(-1,1024)
x = self.dropout(x)
x = self.fc1(x)
if not self.output_features == 1:
x = self.softmax(x)
return x
但是,当我将具有4个图像(全为零)的示例批处理放入这样的模型中时:
x = torch.from_numpy(np.zeros((4,1,256,256))).float()
net = custom_small_CNN(output_features=2,input_channels=1).float()
output = net(x)
输出的形状为torch.Size([16,2])
而不是torch.Size([4,2])
,这正是我想要的,例如ResNet作为输出交付。我想念什么?
谢谢!
解决方法
应用池化层时,它返回(batch_size,2,2,num_filters),因此,当对x = x.view(-1,1024)
进行整形时,其结果是(batch_size * 4,num_filters)为形状。
您应该扁平化或平均池化层的输出,而不是像这样进行重塑。扁平化是这里最常用的。
因此,替换下面的行
x = x.view(-1,1024)
使用
x = nn.Flatten()(x)
将产生正确的最终输出形状
版权声明:本文内容由互联网用户自发贡献,该文观点与技术仅代表作者本人。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如发现本站有涉嫌侵权/违法违规的内容, 请发送邮件至 dio@foxmail.com 举报,一经查实,本站将立刻删除。