用python和torch库实现Unet网络

好的，我可以为您提供一个简单的Unet网络的实现，以下是代码：

import torch
import torch.nn as nn

# 定义编码器
class Encoder(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, batch_norm=True):
        super().__init__()
        if batch_norm:
            self.conv = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding),
                nn.BatchNorm2d(out_channels),
                nn.ReLU(inplace=True),
                nn.Conv2d(out_channels, out_channels, kernel_size, stride, padding),
                nn.BatchNorm2d(out_channels),
                nn.ReLU(inplace=True),
            )
        else:
            self.conv = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding),
                nn.ReLU(inplace=True),
                nn.Conv2d(out_channels, out_channels, kernel_size, stride, padding),
                nn.ReLU(inplace=True),
            )

    def forward(self, x):
        x = self.conv(x)
        return x

# 定义解码器
class Decoder(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, batch_norm=True):
        super().__init__()
        if batch_norm:
            self.conv = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding),
                nn.BatchNorm2d(out_channels),
                nn.ReLU(inplace=True),
                nn.Conv2d(out_channels, out_channels, kernel_size, stride, padding),
                nn.BatchNorm2d(out_channels),
                nn.ReLU(inplace=True),
            )
        else:
            self.conv = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding),
                nn.ReLU(inplace=True),
                nn.Conv2d(out_channels, out_channels, kernel_size, stride, padding),
                nn.ReLU(inplace=True),
            )

    def forward(self, x):
        x = self.conv(x)
        return x

# 定义UNet网络
class UNet(nn.Module):
    def __init__(self, in_channels=3, out_channels=1):
        super().__init__()
        # 编码器部分
        self.enc1 = Encoder(in_channels, 64, batch_norm=False)
        self.pool1 = nn.MaxPool2d(2, 2)
        self.enc2 = Encoder(64, 128)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.enc3 = Encoder(128, 256)
        self.pool3 = nn.MaxPool2d(2, 2)
        self.enc4 = Encoder(256, 512)
        self.pool4 = nn.MaxPool2d(2, 2)

        # 解码器部分
        self.dec4 = Decoder(512, 256)
        self.up4 = nn.ConvTranspose2d(256, 256, 2, stride=2)
        self.dec3 = Decoder(256 + 256, 128)
        self.up3 = nn.ConvTranspose2d(128, 128, 2, stride=2)
        self.dec2 = Decoder(128 + 128, 64)
        self.up2 = nn.ConvTranspose2d(64, 64, 2, stride=2)
        self.dec1 = Decoder(64 + 64, out_channels, batch_norm=False)

    def forward(self, x):
        # 编码器部分
        enc1 = self.enc1(x)
        enc2 = self.enc2(self.pool1(enc1))
        enc3 = self.enc3(self.pool2(enc2))
        enc4 = self.enc4(self.pool3(enc3))

        # 解码器部分
        dec4 = self.dec4(self.pool4(enc4))
        up4 = self.up4(dec4)
        dec3 = self.dec3(torch.cat([up4, enc3], dim=1))
        up3 = self.up3(dec3)
        dec2 = self.dec2(torch.cat([up3, enc2], dim=1))
        up2 = self.up2(dec2)
        dec1 = self.dec1(torch.cat([up2, enc1], dim=1))
        return dec1

这个实现是一个简单的UNet网络，包括编码器和解码器部分，您可以根据您的需求调整网络的深度和宽度。