class CAM_Module(Module): """ Channel attention module""" def __init__(self, in_dim): super(CAM_Module, self).__init__() self.chanel_in = in_dim self.gamma = Parameter(torch.zeros(1)) self.softmax = Softmax(dim=-1) def forward(self,x): """ inputs : x : input feature maps( B X C X H X W) returns : out : attention value + input feature attention: B X C X C """ m_batchsize, C, height, width, channle = x.size() #print(x.size()) proj_query = x.view(m_batchsize, C, -1) proj_key = x.view(m_batchsize, C, -1).permute(0, 2, 1) #形状转换并交换维度 energy = torch.bmm(proj_query, proj_key) energy_new = torch.max(energy, -1, keepdim=True)[0].expand_as(energy)-energy attention = self.softmax(energy_new) proj_value = x.view(m_batchsize, C, -1) out = torch.bmm(attention, proj_value) out = out.view(m_batchsize, C, height, width, channle) # print('out', out.shape) # print('x', x.shape) out = self.gamma*out + x #C*H*W return out 将其改成keras代码

下面是将该 Channel Attention Module 的 PyTorch 代码转换为 Keras 代码的实现：

from keras.layers import Layer, Softmax
import tensorflow as tf

class CAM_Module(Layer):
    """ Channel attention module"""
    def __init__(self, in_dim):
        super(CAM_Module, self).__init__()
        self.chanel_in = in_dim
        self.gamma = self.add_weight(name='gamma', shape=[1], initializer='zeros', trainable=True)
        self.softmax = Softmax(axis=-1)

    def call(self, x):
        """
        inputs :
            x : input feature maps( B X H X W X C)
        returns :
            out : attention value + input feature
            attention: B X C X C
        """
        m_batchsize, height, width, C = x.shape.as_list()
        # print(x.shape)
        proj_query = tf.reshape(x, [-1, height*width, C])
        proj_key = tf.transpose(proj_query, perm=[0, 2, 1])
        energy = tf.matmul(proj_query, proj_key)
        energy_new = tf.expand_dims(tf.reduce_max(energy, axis=-1), axis=-1) - energy
        attention = self.softmax(energy_new)
        proj_value = tf.reshape(x, [-1, height*width, C])
        out = tf.matmul(attention, proj_value)
        out = tf.reshape(out, [-1, height, width, C])
        out = self.gamma * out + x
        return out

注意事项：

1. Keras 中的 Layer 类相当于 PyTorch 中的 Module 类；
2. Keras 中的 add_weight 方法可以添加模型参数，相当于 PyTorch 中的 Parameter 类；
3. Keras 中的 Softmax 层在调用时需要指定 `axis` 参数，相当于 PyTorch 中的 dim 参数；
4. Keras 中的 tf.reshape 方法可以替代 PyTorch 中的 view 方法；
5. Keras 中的 tf.transpose 方法可以替代 PyTorch 中的 permute 方法；
6. Keras 中的 tf.matmul 方法可以替代 PyTorch 中的 torch.bmm 方法。