写一个带有注意力机制的unet网络的py文件，包含：2个类，2个函数，使用nn.module 和 nn.Sequential 和for 循环

下面是一个带有注意力机制的U-Net网络的Python文件示例，包含两个类和两个函数，使用`nn.Module`和`nn.Sequential`以及for循环：

import torch
import torch.nn as nn

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):
        return self.conv(x)


class AttentionBlock(nn.Module):
    def __init__(self, F_g, F_l, F_int):
        super(AttentionBlock, self).__init__()
        self.W_g = nn.Sequential(
            nn.Conv2d(F_g, F_int, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(F_int)
        )

        self.W_x = nn.Sequential(
            nn.Conv2d(F_l, F_int, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(F_int)
        )

        self.psi = nn.Sequential(
            nn.Conv2d(F_int, 1, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(1),
            nn.Sigmoid()
        )

        self.relu = nn.ReLU(inplace=True)

    def forward(self, g, x):
        g1 = self.W_g(g)
        x1 = self.W_x(x)
        psi = self.relu(g1 + x1)
        psi = self.psi(psi)
        out = x * psi
        return out


class UNet(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(UNet, self).__init__()
        self.down1 = DoubleConv(in_channels, 64)
        self.down2 = DoubleConv(64, 128)
        self.up1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
        self.up2 = nn.ConvTranspose2d(64, out_channels, kernel_size=2, stride=2)
        self.attention1 = AttentionBlock(64, 128, 64)
        self.attention2 = AttentionBlock(in_channels, 64, 32)
        self.out_conv = nn.Conv2d(out_channels, out_channels, kernel_size=1)

    def forward(self, x):
        x1 = self.down1(x)
        x2 = self.down2(x1)
        x = self.up1(x2)
        x = self.attention1(x, x2)
        x = self.up2(x + x1)
        x = self.attention2(x, x)
        x = self.out_conv(x)
        return x


# 测试代码
if __name__ == "__main__":
    # 创建一个输入张量
    input_tensor = torch.randn((1, 3, 256, 256))
    
    # 创建UNet模型实例
    model = UNet(in_channels=3, out_channels=3)
    
    # 将输入张量传递给模型进行前向传播
    output_tensor = model(input_tensor)
    
    # 打印输出张量的形状
    print("Output tensor shape:", output_tensor.shape)

在这个示例中，有两个类：`DoubleConv`和`AttentionBlock`，它们都是基于`nn.Module`的自定义模块。`DoubleConv`定义了一个双重卷积层，`AttentionBlock`实现了注意力机制。注意力机制用于增强网络对重要特征的关注。

`UNet`类是基于U-Net架构的模型，它由两个下采样路径（编码器）和两个上采样路径（解码器）组成。在编码器和解码器之间，使用了两个注意力块来引入注意力机制。这些模块和方法通过for循环实现了层次之间的连接。

请注意，这只是一个示例，你可以根据自己的需求修改和扩展这些类和方法。