Unet pytorch

最新推荐文章于 2024-04-01 17:20:14 发布

SoulOSQ

最新推荐文章于 2024-04-01 17:20:14 发布

阅读量407

点赞数

分类专栏： cv 文章标签： pytorch 深度学习 python

本文链接：https://blog.csdn.net/qq_36132590/article/details/126515704

版权

cv 专栏收录该内容

8 篇文章 0 订阅

订阅专栏

import torch
from torch import nn
class double_conv2d_bn(nn.Module):
    def __init__(self,inchannels,outchannels):
        super(double_conv2d_bn,self).__init__()
        self.double_conv_bn = nn.Sequential(
            nn.Conv2d(inchannels,outchannels,kernel_size=3,padding=1,stride=1),
            nn.BatchNorm2d(outchannels),
            nn.ReLU(inplace=True),
            nn.Conv2d(outchannels,outchannels,kernel_size=3,padding=1,stride=1),
            nn.BatchNorm2d(outchannels),
            nn.ReLU(inplace=True)
        )
    def forward(self,inputs):
        return self.double_conv_bn(inputs)
class deconv2d_bn(nn.Module):
    def __init__(self,inchannels,outchannels):
        super(deconv2d_bn,self).__init__()
        if False:
            self.deconv_bn = nn.Sequential(
                nn.Upsample(scale_factor=2),
                nn.Conv2d(inchannels,outchannels,kernel_size=3,padding=1,stride=1),
            )
        else:
            self.deconv_bn = nn.Sequential(
                nn.ConvTranspose2d(inchannels,outchannels,kernel_size=2,stride=2),
                nn.BatchNorm2d(outchannels),
                nn.ReLU(inplace=True)
            )
    def forward(self,inputs):
        return self.deconv_bn(inputs)
class Unet(nn.Module):
    def __init__(self,inchannels=3,outchannels=2):
        super(Unet,self).__init__()
        self.conv1 = double_conv2d_bn(inchannels,64)
        self.maxpool1 = nn.MaxPool2d(kernel_size=2,stride=2)
        self.conv2 = double_conv2d_bn(64,128)
        self.maxpool2 = nn.MaxPool2d(kernel_size=2,stride=2)
        self.conv3 = double_conv2d_bn(128,256)
        self.maxpool3 = nn.MaxPool2d(kernel_size=2,stride=2)
        self.conv4 = double_conv2d_bn(256,512)
        self.maxpool4 = nn.MaxPool2d(kernel_size=2,stride=2)

        self.conv5 = double_conv2d_bn(512,1024)

        self.deconv4 = deconv2d_bn(1024,512)
        self.conv4_ = double_conv2d_bn(1024,512)

        self.deconv3 = deconv2d_bn(512,256)
        self.conv3_ = double_conv2d_bn(512,256)

        self.deconv2 = deconv2d_bn(256,128)
        self.conv2_ = double_conv2d_bn(256,128)

        self.deconv1 = deconv2d_bn(128,64)
        self.conv1_ = double_conv2d_bn(128,64)

        self.out = nn.Conv2d(64,outchannels,kernel_size=1,stride=1)

    def forward(self,inputs):
        conv1 = self.conv1(inputs)
        maxpool1 = self.maxpool1(conv1)

        conv2 = self.conv2(maxpool1)
        maxpool2 = self.maxpool2(conv2)

        conv3 = self.conv3(maxpool2)
        maxpool3 = self.maxpool3(conv3)

        conv4 = self.conv4(maxpool3)
        maxpool4 = self.maxpool4(conv4)

        conv5 = self.conv5(maxpool4)

        deconv4 = self.deconv4(conv5)
        conv4_ = self.conv4_(torch.cat((deconv4,conv4),dim=1))

        deconv3 = self.deconv3(conv4_)
        conv3_ = self.conv3_(torch.cat((deconv3,conv3),dim=1))

        deconv2 = self.deconv2(conv3_)
        conv2_ = self.conv2_(torch.cat((deconv2,conv2),dim=1))

        deconv1 = self.deconv1(conv2_)
        conv1_ = self.conv1_(torch.cat((deconv1,conv1),dim=1))

        out = self.out(conv1_)
        return out

if __name__ == '__main__':
    randnum = torch.rand((4,3,32,32))
    net = Unet()
    out = net(randnum)
    print(out.shape)