pytorch实现segnet_SegNet网络的Pytorch实现

1 importtorch.nn as nn2 importtorch3

4 classconv2DBatchNormRelu(nn.Module):5 def __init__(self,in_channels,out_channels,kernel_size,stride,padding,6 bias=True,dilation=1,is_batchnorm=True):7 super(conv2DBatchNormRelu,self).__init__()8 ifis_batchnorm:9 self.cbr_unit=nn.Sequential(10 nn.Conv2d(in_channels,out_channels,kernel_size=kernel_size,stride=stride,padding=padding,11 bias=bias,dilation=dilation),12 nn.BatchNorm2d(out_channels),13 nn.ReLU(inplace=True),14 )15 else:16 self.cbr_unit=nn.Sequential(17 nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding,18 bias=bias, dilation=dilation),19 nn.ReLU(inplace=True)20 )21

22 defforward(self,inputs):23 outputs=self.cbr_unit(inputs)24 returnoutputs25

26 classsegnetDown2(nn.Module):27 def __init__(self,in_channels,out_channels):28 super(segnetDown2,self).__init__()29 self.conv1=conv2DBatchNormRelu(in_channels,out_channels,kernel_size=3,stride=1,padding=1)30 self.conv2=conv2DBatchNormRelu(out_channels,out_channels,kernel_size=3,stride=1,padding=1)31 self.maxpool_with_argmax=nn.MaxPool2d(kernel_size=2,stride=2,return_indices=True)32

33 defforward(self,inputs):34 outputs=self.conv1(inputs)35 outputs=self.conv2(outputs)36 unpooled_shape=outputs.size()37 outputs,indices=self.maxpool_with_argmax(outputs)38 returnoutputs,indices,unpooled_shape39

40 classsegnetDown3(nn.Module):41 def __init__(self,in_channels,out_channels):42 super(segnetDown3,self).__init__()43 self.conv1=conv2DBatchNormRelu(in_channels,out_channels,kernel_size=3,stride=1,padding=1)44 self.conv2=conv2DBatchNormRelu(out_channels,out_channels,kernel_size=3,stride=1,padding=1)45 self.conv3=conv2DBatchNormRelu(out_channels,out_channels,kernel_size=3,stride=1,padding=1)46 self.maxpool_with_argmax=nn.MaxPool2d(kernel_size=2,stride=2,return_indices=True)47

48 defforward(self,inputs):49 outputs=self.conv1(inputs)50 outputs=self.conv2(outputs)51 outputs=self.conv3(outputs)52 unpooled_shape=outputs.size()53 outputs,indices=self.maxpool_with_argmax(outputs)54 returnoutputs,indices,unpooled_shape55

56

57 classsegnetUp2(nn.Module):58 def __init__(self,in_channels,out_channels):59 super(segnetUp2,self).__init__()60 self.unpool=nn.MaxUnpool2d(2,2)61 self.conv1=conv2DBatchNormRelu(in_channels,out_channels,kernel_size=3,stride=1,padding=1)62 self.conv2=conv2DBatchNormRelu(out_channels,out_channels,kernel_size=3,stride=1,padding=1)63

64 defforward(self,inputs,indices,output_shape):65 outputs=self.unpool(inputs,indices=indices,output_size=output_shape)66 outputs=self.conv1(outputs)67 outputs=self.conv2(outputs)68 returnoutputs69

70 classsegnetUp3(nn.Module):71 def __init__(self,in_channels,out_channels):72 super(segnetUp3,self).__init__()73 self.unpool=nn.MaxUnpool2d(2,2)74 self.conv1=conv2DBatchNormRelu(in_channels,out_channels,kernel_size=3,stride=1,padding=1)75 self.conv2=conv2DBatchNormRelu(out_channels,out_channels,kernel_size=3,stride=1,padding=1)76 self.conv3=conv2DBatchNormRelu(out_channels,out_channels,kernel_size=3,stride=1,padding=1)77

78 defforward(self,inputs,indices,output_shape):79 outputs=self.unpool(inputs,indices=indices,output_size=output_shape)80 outputs=self.conv1(outputs)81 outputs=self.conv2(outputs)82 outputs=self.conv3(outputs)83 returnoutputs84

85 classsegnet(nn.Module):86 def __init__(self,in_channels=3,num_classes=21):87 super(segnet,self).__init__()88 self.down1=segnetDown2(in_channels=in_channels,out_channels=64)89 self.down2=segnetDown2(64,128)90 self.down3=segnetDown3(128,256)91 self.down4=segnetDown3(256,512)92 self.down5=segnetDown3(512,512)93

94 self.up5=segnetUp3(512,512)95 self.up4=segnetUp3(512,256)96 self.up3=segnetUp3(256,128)97 self.up2=segnetUp2(128,64)98 self.up1=segnetUp2(64,64)99 self.finconv=conv2DBatchNormRelu(64,num_classes,3,1,1)100

101 defforward(self,inputs):102 down1,indices_1,unpool_shape1=self.down1(inputs)103 down2,indices_2,unpool_shape2=self.down2(down1)104 down3,indices_3,unpool_shape3=self.down3(down2)105 down4,indices_4,unpool_shape4=self.down4(down3)106 down5,indices_5,unpool_shape5=self.down5(down4)107

108 up5=self.up5(down5,indices=indices_5,output_shape=unpool_shape5)109 up4=self.up4(up5,indices=indices_4,output_shape=unpool_shape4)110 up3=self.up3(up4,indices=indices_3,output_shape=unpool_shape3)111 up2=self.up2(up3,indices=indices_2,output_shape=unpool_shape2)112 up1=self.up1(up2,indices=indices_1,output_shape=unpool_shape1)113 outputs=self.finconv(up1)114

115 returnoutputs116

117 if __name__=="__main__":118 inputs=torch.ones(1,3,224,224)119 model=segnet()120 print(model(inputs).size())121 print(model)