Unet应用完整代码举例

以下是一个使用PyTorch实现的Unet网络的完整代码示例：

import torch
import torch.nn as nn

# 定义Unet网络模型
class Unet(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(Unet, self).__init__()
        
        # 定义编码器部分
        self.down1 = self.contracting_block(in_channels, 64)
        self.down2 = self.contracting_block(64, 128)
        self.down3 = self.contracting_block(128, 256)
        self.down4 = self.contracting_block(256, 512)
        
        # 定义解码器部分
        self.up1 = self.expanding_block(512, 256)
        self.up2 = self.expanding_block(256, 128)
        self.up3 = self.expanding_block(128, 64)
        
        # 定义最后的卷积层
        self.out = nn.Sequential(
            nn.Conv2d(64, out_channels, kernel_size=1),
            nn.Sigmoid()
        )
        
    def contracting_block(self, in_channels, out_channels):
        block = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        return block
    
    def expanding_block(self, in_channels, out_channels):
        block = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(out_channels, in_channels // 2, kernel_size=2, stride=2)
        )
        return block
    
    def forward(self, x):
        # 编码器部分
        x1 = self.down1(x)
        x2 = self.down2(x1)
        x3 = self.down3(x2)
        x4 = self.down4(x3)
        
        # 解码器部分
        x = self.up1(x4)
        x = self.up2(torch.cat([x, x3], dim=1))
        x = self.up3(torch.cat([x, x2], dim=1))
        
        # 输出层
        x = self.out(torch.cat([x, x1], dim=1))
        return x


# 创建Unet模型实例
unet = Unet(in_channels=3, out_channels=1)

# 输入数据
inputs = torch.randn(1, 3, 256, 256)

# 前向传播
outputs = unet(inputs)
print(outputs.shape)

这个代码示例定义了一个Unet网络的模型，包括编码器部分、解码器部分和最后的卷积层。你可以根据自己的需求调整输入通道数和输出通道数。在前向传播过程中，我们输入了一个大小为1x3x256x256的张量，得到了一个大小为1x1x256x256的输出。