基于python实现resnet_python - 基于ResNet的UNet的Pytorch实现 - 堆栈内存溢出

我想实现基于ResNet的UNet进行细分(无需预先培训)。 我已经使用Keras提到了这个实现，但是我的项目已经使用PyTorch实现了，我不确定我是否做得正确。

我的PyTorch实现(不确定我是否正确...) 任何建议将不胜感激 。

def double_conv(in_channels, out_channels):

return nn.Sequential(

nn.Conv2d(in_channels, out_channels, 3, padding=1),

nn.BatchNorm2d(out_channels),

nn.ReLU(inplace=True),

nn.Conv2d(out_channels, out_channels, 3, padding=1),

nn.BatchNorm2d(out_channels),

nn.ReLU(inplace=True)

)

def root_block(in_channels, out_channels):

return nn.Sequential(

nn.Conv2d(in_channels, out_channels, 3, padding=1),

nn.BatchNorm2d(out_channels),

nn.ReLU(inplace=True),

nn.Conv2d(out_channels, out_channels, 3, padding=1),

nn.BatchNorm2d(out_channels),

)

# Define the UNet architecture

class ResNetUNet(nn.Module):

def __init__(self, n_class):

super().__init__()

self.dconv_down1 = double_conv(3, 64)

self.dconv_down11 = root_block(64, 64)

self.dconv_down2 = double_conv(64, 128)

self.dconv_down21 = root_block(128, 128)

self.dconv_down3 = double_conv(128, 256)

self.dconv_down31 = root_block(256, 256)

self.dconv_down4 = double_conv(256, 512)

self.dconv_down41 = root_block(512, 512)

self.maxpool = nn.MaxPool2d(2)

self.relu = nn.ReLU(inplace=True)

self.dconv_up3 = double_conv(256 + 512, 256)

self.dconv_up31 = root_block(256, 256)

self.dconv_up2 = double_conv(128 + 256, 128)

self.dconv_up21 = root_block(128, 128)

self.dconv_up1 = double_conv