Pytorch搭建U-Net网络

U-Net: Convolutional Networks for Biomedical Image Segmentation

import torch.nn as nn
import torch
from torch import autograd
from torchsummary import summary

class DoubleConv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(DoubleConv, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=0),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, 3, padding=0),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )

    def forward(self, input):
        return self.conv(input)

class Unet(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(Unet, self).__init__()
        self.conv1 = DoubleConv(in_ch, 64)
        self.pool1 = nn.MaxPool2d(2)
        self.conv2 = DoubleConv(64, 128)
        self.pool2 = nn.MaxPool2d(2)
        self.conv3 = DoubleConv(128, 256)
        self.pool3 = nn.MaxPool2d(2)
        self.conv4 = DoubleConv(256, 512)
        self.pool4 = nn.MaxPool2d(2)
        self.conv5 = DoubleConv(512, 1024)
        # 逆卷积，也可以使用上采样
        self.up6 = nn.ConvTranspose2d(1024, 512, 2, stride=2)
        self.conv6 = DoubleConv(1024, 512)
        self.up7 = nn.ConvTranspose2d(512, 256, 2, stride=2)
        self.conv7 = DoubleConv(512, 256)
        self.up8 = nn.ConvTranspose2d(256, 128, 2, stride=2)
        self.conv8 = DoubleConv(256, 128)
        self.up9 = nn.ConvTranspose2d(128, 64, 2, stride=2)
        self.conv9 = DoubleConv(128, 64)
        self.conv10 = nn.Conv2d(64, out_ch, 1)

    def forward(self, x):
        c1 = self.conv1(x)
        crop1 = c1[:,:,88:480,88:480]
        p1 = self.pool1(c1)
        c2 = self.conv2(p1)
        crop2 = c2[:,:,40:240,40:240]
        p2 = self.pool2(c2)
        c3 = self.conv3(p2)
        crop3 = c3[:,:,16:120,16:120]
        p3 = self.pool3(c3)
        c4 = self.conv4(p3)
        crop4 = c4[:,:,4:60,4:60]
        p4 = self.pool4(c4)
        c5 = self.conv5(p4)
        up_6 = self.up6(c5)
        merge6 = torch.cat([up_6, crop4], dim=1)
        c6 = self.conv6(merge6)
        up_7 = self.up7(c6)
        merge7 = torch.cat([up_7, crop3], dim=1)
        c7 = self.conv7(merge7)
        up_8 = self.up8(c7)
        merge8 = torch.cat([up_8, crop2], dim=1)
        c8 = self.conv8(merge8)
        up_9 = self.up9(c8)
        merge9 = torch.cat([up_9, crop1], dim=1)
        c9 = self.conv9(merge9)
        c10 = self.conv10(c9)
        out = nn.Sigmoid()(c10)
        return out


if __name__=="__main__":
    test_input=torch.rand(1, 1, 572, 572)
    model=Unet(in_ch=1, out_ch=2)
    summary(model, (1,572,572))
    ouput=model(test_input)
    print(ouput.size())

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1         [-1, 64, 570, 570]             640
       BatchNorm2d-2         [-1, 64, 570, 570]             128
              ReLU-3         [-1, 64, 570, 570]               0
            Conv2d-4         [-1, 64, 568, 568]          36,928
       BatchNorm2d-5         [-1, 64, 568, 568]             128
              ReLU-6         [-1, 64, 568, 568]               0
        DoubleConv-7         [-1, 64, 568, 568]               0
         MaxPool2d-8         [-1, 64, 284, 284]               0
            Conv2d-9        [-1, 128, 282, 282]          73,856
      BatchNorm2d-10        [-1, 128, 282, 282]             256
             ReLU-11        [-1, 128, 282, 282]               0
           Conv2d-12        [-1, 128, 280, 280]         147,584
      BatchNorm2d-13        [-1, 128, 280, 280]             256
             ReLU-14        [-1, 128, 280, 280]               0
       DoubleConv-15        [-1, 128, 280, 280]               0
        MaxPool2d-16        [-1, 128, 140, 140]               0
           Conv2d-17        [-1, 256, 138, 138]         295,168
      BatchNorm2d-18        [-1, 256, 138, 138]             512
             ReLU-19        [-1, 256, 138, 138]               0
           Conv2d-20        [-1, 256, 136, 136]         590,080
      BatchNorm2d-21        [-1, 256, 136, 136]             512
             ReLU-22        [-1, 256, 136, 136]               0
       DoubleConv-23        [-1, 256, 136, 136]               0
        MaxPool2d-24          [-1, 256, 68, 68]               0
           Conv2d-25          [-1, 512, 66, 66]       1,180,160
      BatchNorm2d-26          [-1, 512, 66, 66]           1,024
             ReLU-27          [-1, 512, 66, 66]               0
           Conv2d-28          [-1, 512, 64, 64]       2,359,808
      BatchNorm2d-29          [-1, 512, 64, 64]           1,024
             ReLU-30          [-1, 512, 64, 64]               0
       DoubleConv-31          [-1, 512, 64, 64]               0
        MaxPool2d-32          [-1, 512, 32, 32]               0
           Conv2d-33         [-1, 1024, 30, 30]       4,719,616
      BatchNorm2d-34         [-1, 1024, 30, 30]           2,048
             ReLU-35         [-1, 1024, 30, 30]               0
           Conv2d-36         [-1, 1024, 28, 28]       9,438,208
      BatchNorm2d-37         [-1, 1024, 28, 28]           2,048
             ReLU-38         [-1, 1024, 28, 28]               0
       DoubleConv-39         [-1, 1024, 28, 28]               0
  ConvTranspose2d-40          [-1, 512, 56, 56]       2,097,664
           Conv2d-41          [-1, 512, 54, 54]       4,719,104
      BatchNorm2d-42          [-1, 512, 54, 54]           1,024
             ReLU-43          [-1, 512, 54, 54]               0
           Conv2d-44          [-1, 512, 52, 52]       2,359,808
      BatchNorm2d-45          [-1, 512, 52, 52]           1,024
             ReLU-46          [-1, 512, 52, 52]               0
       DoubleConv-47          [-1, 512, 52, 52]               0
  ConvTranspose2d-48        [-1, 256, 104, 104]         524,544
           Conv2d-49        [-1, 256, 102, 102]       1,179,904
      BatchNorm2d-50        [-1, 256, 102, 102]             512
             ReLU-51        [-1, 256, 102, 102]               0
           Conv2d-52        [-1, 256, 100, 100]         590,080
      BatchNorm2d-53        [-1, 256, 100, 100]             512
             ReLU-54        [-1, 256, 100, 100]               0
       DoubleConv-55        [-1, 256, 100, 100]               0
  ConvTranspose2d-56        [-1, 128, 200, 200]         131,200
           Conv2d-57        [-1, 128, 198, 198]         295,040
      BatchNorm2d-58        [-1, 128, 198, 198]             256
             ReLU-59        [-1, 128, 198, 198]               0
           Conv2d-60        [-1, 128, 196, 196]         147,584
      BatchNorm2d-61        [-1, 128, 196, 196]             256
             ReLU-62        [-1, 128, 196, 196]               0
       DoubleConv-63        [-1, 128, 196, 196]               0
  ConvTranspose2d-64         [-1, 64, 392, 392]          32,832
           Conv2d-65         [-1, 64, 390, 390]          73,792
      BatchNorm2d-66         [-1, 64, 390, 390]             128
             ReLU-67         [-1, 64, 390, 390]               0
           Conv2d-68         [-1, 64, 388, 388]          36,928
      BatchNorm2d-69         [-1, 64, 388, 388]             128
             ReLU-70         [-1, 64, 388, 388]               0
       DoubleConv-71         [-1, 64, 388, 388]               0
           Conv2d-72          [-1, 2, 388, 388]             130
================================================================
Total params: 31,042,434
Trainable params: 31,042,434
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 1.25
Forward/backward pass size (MB): 3280.59
Params size (MB): 118.42
Estimated Total Size (MB): 3400.26
----------------------------------------------------------------
torch.Size([1, 2, 388, 388])