class UNetEx(nn.Layer): def __init__(self, in_channels, out_channels, kernel_size=3, filters=[16, 32, 64], layers=3, weight_norm=True, batch_norm=True, activation=nn.ReLU, final_activation=None): super().__init__() assert len(filters) > 0 self.final_activation = final_activation self.encoder = create_encoder(in_channels, filters, kernel_size, weight_norm, batch_norm, activation, layers) decoders = [] for i in range(out_channels): decoders.append(create_decoder(1, filters, kernel_size, weight_norm, batch_norm, activation, layers)) self.decoders = nn.Sequential(*decoders) def encode(self, x): tensors = [] indices = [] sizes = [] for encoder in self.encoder: x = encoder(x) sizes.append(x.shape) tensors.append(x) x, ind = F.max_pool2d(x, 2, 2, return_mask=True) indices.append(ind) return x, tensors, indices, sizes def decode(self, _x, _tensors, _indices, _sizes): y = [] for _decoder in self.decoders: x = _x tensors = _tensors[:] indices = _indices[:] sizes = _sizes[:] for decoder in _decoder: tensor = tensors.pop() size = sizes.pop() ind = indices.pop() # 反池化操作，为上采样 x = F.max_unpool2d(x, ind, 2, 2, output_size=size) x = paddle.concat([tensor, x], axis=1) x = decoder(x) y.append(x) return paddle.concat(y, axis=1) def forward(self, x): x, tensors, indices, sizes = self.encode(x) x = self.decode(x, tensors, indices, sizes) if self.final_activation is not None: x = self.final_activation(x) return x 不修改上述神经网络的encoder和decoder的生成方式，用嘴少量的代码实现attention机制，在上述代码里修改。

可以在UNetEx类中加入一个Attention模块，然后在decode函数中使用该Attention模块。具体实现如下：

class Attention(nn.Layer):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv1 = nn.Conv2D(in_channels, out_channels, kernel_size=1)
        self.conv2 = nn.Conv2D(out_channels, out_channels, kernel_size=1)

    def forward(self, x, y):
        x = self.conv1(x)
        y = self.conv2(y)
        z = paddle.add(x, y)
        z = nn.functional.sigmoid(z)
        z = paddle.multiply(x, z)
        return z

class UNetEx(nn.Layer):
    def __init__(self, in_channels, out_channels, kernel_size=3, filters=[16, 32, 64], layers=3,
                 weight_norm=True, batch_norm=True, activation=nn.ReLU, final_activation=None):
        super().__init__()
        assert len(filters) > 0
        self.final_activation = final_activation
        self.encoder = create_encoder(in_channels, filters, kernel_size, weight_norm, batch_norm, activation, layers)
        self.attention = Attention(filters[-1], filters[-1])
        decoders = []
        for i in range(out_channels):
            decoders.append(create_decoder(1, filters, kernel_size, weight_norm, batch_norm, activation, layers))
        self.decoders = nn.Sequential(*decoders)

    def encode(self, x):
        tensors = []
        indices = []
        sizes = []
        for encoder in self.encoder:
            x = encoder(x)
            sizes.append(x.shape)
            tensors.append(x)
            x, ind = F.max_pool2d(x, 2, 2, return_mask=True)
            indices.append(ind)
        return x, tensors, indices, sizes

    def decode(self, _x, _tensors, _indices, _sizes):
        y = []
        for _decoder in self.decoders:
            x = _x
            tensors = _tensors[:]
            indices = _indices[:]
            sizes = _sizes[:]
            for decoder in _decoder:
                tensor = tensors.pop()
                size = sizes.pop()
                ind = indices.pop()
                # 反池化操作，为上采样
                x = F.max_unpool2d(x, ind, 2, 2, output_size=size)
                x = self.attention(tensor, x)  # 使用Attention模块
                x = decoder(x)
            y.append(x)
        return paddle.concat(y, axis=1)

    def forward(self, x):
        x, tensors, indices, sizes = self.encode(x)
        x = self.decode(x, tensors, indices, sizes)
        if self.final_activation is not None:
            x = self.final_activation(x)
        return x

在该代码中，我们增加了一个Attention类，它接收两个特征图，通过两个1x1卷积层将它们映射到同一维度，然后将它们相加并通过sigmoid函数归一化，最后将第一个特征图与归一化后的结果相乘得到注意力加权后的特征图。在UNetEx类中，我们在decoder函数中使用了Attention类，并将encoder中的最后一层特征图与decoder中的每一层特征图进行注意力加权。这样就实现了在UNetEx中加入Attention机制。