用代码向我展示UNet网络结构

我可以用下面的代码向您展示UNet网络结构：class UNet(nn.Module):
def __init__(self, in_channels, out_channels):
super().__init__()

self.down_conv1 = nn.Conv2d(in_channels, 64, 3, padding=1)
self.down_conv2 = nn.Conv2d(64, 64, 3, padding=1)
self.down_maxpool = nn.MaxPool2d(2)

self.up_conv1 = nn.Conv2d(64, 64, 3, padding=1)
self.up_conv2 = nn.Conv2d(64, out_channels, 3, padding=1)

def forward(self, x):
# Downsample
x1 = self.down_maxpool(F.relu(self.down_conv1(x)))
x2 = self.down_maxpool(F.relu(self.down_conv2(x1)))

# Upsample
x = F.interpolate(x2, scale_factor=2, mode='bilinear', align_corners=True)
x = F.relu(self.up_conv1(x))
x = self.up_conv2(x)

return x