u-net代码(多类别训练,pytorch)

U-Net是一个经典的语义分割模型，常用于医学图像处理。相比传统的卷积神经网络，U-Net在网络结构上采用了类似于自编码器的对称结构，在上采样的过程中使用了Skip Connection技术，能够更好的保留图像中物体之间的空间关系。

在多类别训练中，U-Net模型需要根据实际情况设计相应的输出层结构，实现多标签的分类。在Pytorch中，可以使用nn.Module构建U-Net模型的各个模块，具体实现：

1. 定义U-Net模型：

class UNet(nn.Module):
    def __init__(self, n_channels, n_classes):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes

        self.inc = DoubleConv(n_channels, 64)
        self.down1 = Down(64, 128)
        self.down2 = Down(128, 256)
        self.down3 = Down(256, 512)
        self.down4 = Down(512, 512)
        self.up1 = Up(1024, 256)
        self.up2 = Up(512, 128)
        self.up3 = Up(256, 64)
        self.up4 = Up(128, 64)
        self.outc = nn.Conv2d(64, n_classes, 1)

    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.outc(x)
        return logits

其中， n_channels为输入图像通道数，n_classes为输出类别数，inc代表输入的首个卷积层，down代表下采样过程中的卷积层，up代表上采样过程中的卷积层，outc代表输出的类别数。

2. 定义DoubleConv层：

class DoubleConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(DoubleConv, self).__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):
        x = self.conv(x)
        return x

这里使用了两个卷积层，分别之间使用了BatchNorm和ReLU函数做归一化和激活函数。

3. 定义Down和Up层：

class Down(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(Down, self).__init__()
        self.mpconv = nn.Sequential(
            nn.MaxPool2d(2),
            DoubleConv(in_channels, out_channels)
        )

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

class Up(nn.Module):
    def __init__(self, in_channels, out_channels, bilinear=True):
        super(Up, self).__init__()
        if bilinear:
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        else:
            self.up = nn.ConvTranspose2d(in_channels//2, in_channels//2, kernel_size=2, stride=2)

        self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1, x2):
        x1 = self.up(x1)
        diffX = x2.size()[2] - x1.size()[2]
        diffY = x2.size()[3] - x1.size()[3]
        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2])
        x = torch.cat([x2, x1], dim=1)
        x = self.conv(x)
        return x

Down层中使用了MaxPool下采样，而Up层中采用了上采样的方式还原分辨率，并使用了torch.cat()函数将两组下采样和上采样的两层特征合并。

4. 模型训练：

在训练中，需要定义损失函数和优化器，并将图像数据和标签数据按batch传入模型中计算loss和更新梯度。

model = UNet(n_channels=1, n_classes=2).to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=lr)

for epoch in range(epochs):
    running_loss = 0.0
    for i, data in enumerate(train_loader, 0):
        inputs, labels = data[0].to(device), data[1].to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()

    if epoch % 10 == 9:
        print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 10))
        running_loss = 0.0

其中，使用了CrossEntropyLoss函数来计算多分类的损失值。在更新过程中使用了Adam优化器来更新梯度。

以上是U-Net多类别训练的Pytorch实现方法，需要注意的是，在实际训练过程中不同场景下需要进行相应的调整和改进，以达到更好的训练效果。