如何用基于resnet的Unet进行图像分割基于Pytorch-0.5版本--688IT编程网

如何⽤基于resnet的Unet进⾏图像分割基于Pytorch-0.5版本1、关于Unet

Unet主要⽤于医学图像的很多论⽂中，以及Kaggle竞赛和⼀些其他竞赛中“少类别”的图像分割。从我做实验的经验来说，像VOC这种类别⽐较多的分割任务，不容易收敛，效果较为差。

2、Resnet34

我们的encode部分选择resnet34，decode部分为每⼀个block制作三层卷积，其中每个的第⼆层为upsample（bilinear/deconv）

#基本的block

class DecoderBlock(nn.Module):

def __init__(self,

in_channels=512,

n_filters=256,

kernel_size=3,

is_deconv=False,

super().__init__()

if kernel_size == 3:

conv_padding = 1

elif kernel_size == 1:

conv_padding = 0

# B, C, H, W -> B, C/4, H, W

in_channels // 4,

kernel_size,

padding=1,bias=False)

<1 = nn.BatchNorm2d(in_channels // 4)

# B, C/4, H, W -> B, C/4, H, W

if is_deconv == True:

self.deconv2 = nn.ConvTranspose2d(in_channels // 4,

in_channels // 4,

stride=2,

padding=1,

output_padding=conv_padding,bias=False)

else:

self.deconv2 = nn.Upsample(scale_factor=2,**up_kwargs)

<2 = nn.BatchNorm2d(in_channels // 4)

# B, C/4, H, W -> B, C, H, W

n_filters,

kernel_size,

padding=conv_padding,bias=False)

<3 = nn.BatchNorm2d(n_filters)

def forward(self, x):

x = v1(x)

x = 1(x)

x = lu1(x)

x = self.deconv2(x)

x = 2(x)

x = lu2(x)

x = v3(x)

x = 3(x)

x = lu3(x)

return x

class BasicConv2d(nn.Module):

def __init__(self, in_planes, out_planes, kernel_size, stride, padding=0):

super(BasicConv2d, self).__init__()

kernel_size=kernel_size, stride=stride,

padding=padding, bias=False) # verify bias false

self.bn = nn.BatchNorm2d(out_planes,

eps=0.001, # value found in tensorflow

momentum=0.1, # default pytorch value

affine=True)

def forward(self, x):

x = v(x)

x = self.bn(x)

x = lu(x)

return x

接下来，就是模型代码了。

class ResNet34Unet(nn.Module):

def __init__(self,

num_classes,

num_channels=3,

is_deconv=False,

decoder_kernel_size=3,

super().__init__()

self.std = (0.229, 0.224, 0.225)

filters = [64, 128, 256, 512]

resnet = snet34(pretrained=True)tensorflow版本选择

self.base_size=512

self._up_kwargs={'mode': 'bilinear', 'align_corners': True}

# self.firstconv = v1

# assert num_channels == 3, "num channels not used now. to use changle first conv layer to support num channels other then 3" # try to use 8-channels as first input

if num_channels == 3:

self.firstconv = v1

else:

self.firstconv = nn.Conv2d(num_channels, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3),bias=False)

self.firstconv = nn.Conv2d(num_channels, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3),bias=Fa

lse)

self.firstbn = resnet.bn1

self.firstrelu = lu

self.firstmaxpool = resnet.maxpool

# Decoder

< = DecoderBlock(in_channels=filters[3],

n_filters=filters[3],