基于PyTorch的UNet预训练模型

你可以使用torchvision内置的预训练模型来实现UNet模型，例如ResNet、VGG等。这些模型已经在ImageNet上进行了预训练，可以直接用于图像分割任务。

下面是一个使用ResNet50作为编码器的UNet模型的示例代码：

import torch
import torch.nn as nn
import torchvision.models as models

class DoubleConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, 3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

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

class UNet(nn.Module):
    def __init__(self, n_classes):
        super().__init__()
        self.encoder = models.resnet50(pretrained=True)
        self.pool = nn.MaxPool2d(2, 2)

        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)

        self.conv0 = DoubleConv(3, 64)
        self.conv1 = DoubleConv(64, 128)
        self.conv2 = DoubleConv(128, 256)
        self.conv3 = DoubleConv(256, 512)
        self.conv4 = DoubleConv(512, 1024)

        self.center = DoubleConv(2048, 1024)

        self.dec4 = DoubleConv(1024 + 512, 512)
        self.dec3 = DoubleConv(512 + 256, 256)
        self.dec2 = DoubleConv(256 + 128, 128)
        self.dec1 = DoubleConv(128 + 64, 64)

        self.final = nn.Conv2d(64, n_classes, 1)

    def forward(self, x):
        conv0 = self.conv0(x)
        conv1 = self.conv1(self.pool(conv0))
        conv2 = self.conv2(self.pool(conv1))
        conv3 = self.conv3(self.pool(conv2))
        conv4 = self.conv4(self.pool(conv3))

        center = self.center(torch.cat([conv4, self.encoder.conv1(conv3)], dim=1))

        dec4 = self.dec4(torch.cat([center, conv4], dim=1))
        dec3 = self.dec3(torch.cat([self.up(dec4), conv3], dim=1))
        dec2 = self.dec2(torch.cat([self.up(dec3), conv2], dim=1))
        dec1 = self.dec1(torch.cat([self.up(dec2), conv1], dim=1))

        return self.final(dec1)

在这个模型中，首先使用预训练的ResNet50作为编码器，然后添加了几个上采样和下采样的模块，最后通过一个1x1卷积层输出预测结果。这个模型可以用于二分类任务，如果要进行多分类任务，只需要将最后一个卷积层的输出通道数改为类别数即可。

在实际使用中，可以根据自己的任务需求进行调整和修改。