最近修改了两个神经网络模型(3D Res-UNet和3D V-Net):原来输入张量的大小为(1, 1, 16, 512, 512)
输出张量的大小也为(1, 1, 16, 512, 512)
改为 输入(1, 2, 5, 512, 512)
输出(1, 2, 5, 512, 512)
由于3D Res-UNet和3D V-Net都是U-Net网络的变体,所以都会涉及到下采样和上采样的问题。
由于原来是一次输入16张512512大小的图片,16为偶数,所以经过下采样和上采样后能够同样得到16张。在输入和当前阶段的输出做相加操作时也不会因为维度不同而报错。
但是改为一次输出5张512512大小的图片时,5为奇数,经过下采样后到1而上采样得到的却只能是偶数(2 4 8 16…)回不到5了。
self.down_conv1 = nn.Sequential(
nn.Conv3d(16, 32, 2, 2),
nn.PReLU(32)
)
self.down_conv2 = nn.Sequential(
nn.Conv3d(32, 64, 2, 2),
nn.PReLU(64)
)
self.down_conv3 = nn.Sequential(
nn.Conv3d(64, 128, 2, 2),
nn.PReLU(128)
)
self.down_conv4 = nn.Sequential(
nn.Conv3d(128, 256, 3, 1, padding=1),
nn.PReLU(256)
)
self.up_conv2 = nn.Sequential(
nn.ConvTranspose3d(256, 128, 2, 2), # 进行转置卷积
nn.PReLU(128)
)
self.up_conv3 = nn.Sequential(
nn.ConvTranspose3d(128, 64, 2, 2),
nn.PReLU(64)
)
self.up_conv4 = nn.Sequential(
nn.ConvTranspose3d(64, 32, 2, 2),
nn.PReLU(32)
)
只需要将 在下采样和上采样过程中卷积核的大小步长由size=2, stride =2
改为size=(1, 2, 2), stride =(1, 2, 2)
。这样由于下采样和上采样的过程中对于图片张数的卷积核和步长均为1,无论一次性输出多少的张的图片 网络的下采样和上采样均不会受到影响
self.down_conv1 = nn.Sequential(
nn.Conv3d(16, 32, (1, 2, 2), (1, 2, 2)),
nn.PReLU(32)
)
self.down_conv2 = nn.Sequential(
nn.Conv3d(32, 64, (1, 2, 2), (1, 2, 2)),
nn.PReLU(64)
)
self.down_conv3 = nn.Sequential(
nn.Conv3d(64, 128, (1, 2, 2), (1, 2, 2)),
nn.PReLU(128)
)
self.down_conv4 = nn.Sequential(
nn.Conv3d(128, 256, 3, 1, padding=1),
nn.PReLU(256)
)
self.up_conv1 = nn.Sequential(
nn.ConvTranspose3d(256, 128, (1, 2, 2), (1, 2, 2)), # 进行转置卷积
nn.PReLU(128)
)
self.up_conv2 = nn.Sequential(
nn.ConvTranspose3d(128, 64, (1, 2, 2), (1, 2, 2)),
nn.PReLU(64)
)
self.up_conv3 = nn.Sequential(
nn.ConvTranspose3d(64, 32, (1, 2, 2), (1, 2, 2)),
nn.PReLU(32)
)
if __name__ == '__main__':
X = torch.randn(size = (1, 2, 5, 512, 512))
print("输入数据的大小为:", X.shape)
net = ResUNet(in_channels=2, out_channels=2, training=True)
# common.print_network(net)
output = net(X)
print("输出结果的大小为:", output.shape)