用pytorch打印UNET

以下是使用 PyTorch 打印 U-Net 模型的示例代码：

import torch
import torch.nn as nn

class UNet(nn.Module):
    def __init__(self):
        super().__init__()

        self.down1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True)
        )

        self.down2 = nn.Sequential(
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True)
        )

        self.up1 = nn.Sequential(
            nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True)
        )

        self.output = nn.Conv2d(64, 3, kernel_size=1)

    def forward(self, x):
        x = self.down1(x)
        x = self.down2(x)
        x = self.up1(x)
        x = self.output(x)

        return x

model = UNet()
print(model)

这段代码定义了一个简单的 U-Net 模型，并使用 `print(model)` 将其打印出来。该模型有两个下采样层和一个上采样层，同时对最终的输出进行了卷积。这个示例代码仅供参考，实际的 U-Net 模型可能会更加复杂。