写一段基于pytorch的unet神经网络求解回归任务的python代码

好的，下面是基于PyTorch的UNet神经网络求解回归任务的Python代码示例。

import torch
import torch.nn as nn
import torch.optim as optim

class UNet(nn.Module):
    def __init__(self):
        super(UNet, self).__init__()
        
        self.down1 = DoubleConv(1, 64)
        self.down2 = DoubleConv(64, 128)
        self.down3 = DoubleConv(128, 256)
        self.down4 = DoubleConv(256, 512)
        self.up1 = Up(512, 256) 
        self.up2 = Up(256, 128) 
        self.up3 = Up(128, 64) 
        self.up4 = Up(64, 32)
        self.out = nn.Sequential(nn.Conv2d(32, 1, kernel_size=1), nn.Sigmoid())

        self.maxpool = nn.MaxPool2d(2, stride=2)
        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        
    def forward(self, x):
        x1 = self.down1(x)
        x2 = self.maxpool(x1)
        x3 = self.down2(x2)
        x4 = self.maxpool(x3)  
        x5 = self.down3(x4)
        x6 = self.maxpool(x5)  
        x7 = self.down4(x6)
        x = self.upsample(x7)
        x8 = torch.cat([x, x5], dim=1)
        x8 = self.up1(x8)
        x = self.upsample(x8)
        x9 = torch.cat([x, x3], dim=1)
        x9 = self.up2(x9)
        x = self.upsample(x9)
        x10 = torch.cat([x, x1], dim=1)
        x10 = self.up3(x10)
        x = self.upsample(x10)
        x = self.up4(x)
        x = self.out(x)
        return x
        
        
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, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
        )

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

class Up(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(Up, self).__init__()
        self.conv = DoubleConv(in_channels + out_channels, out_channels)
        
    def forward(self, x1, x2):
        x1 = self.upsample(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

model = UNet()
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 数据准备
...
# 训练过程
for epoch in range(num_epochs):
    for data, target in train_loader:
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()