从 keras 到 pytorch 的自定义层 [英] Custom layer from keras to pytorch

问题描述

来自 TensorFlow 背景，我正在尝试将自定义层的一段代码从 Keras 转换为 PyTorch.

Coming from TensorFlow background, I am trying to convert a snippet of code of the custom layer from Keras to PyTorch.

Keras 中的自定义层如下所示:

The custom layer in Keras looks like this:

class Attention_module(tf.keras.layers.Layer):
    def __init__(self, class_num):
        super(Attention_module,self).__init__(class_num)
        self.class_num = class_num
        self.Ws = None

    def build(self, input_shape):
        embedding_length = int(input_shape[2])

        self.Ws = self.add_weight(shape=(self.class_num, embedding_length),
                                  initializer=tf.keras.initializers.get('glorot_uniform'), trainable=True)

        super(Attention_module, self).build(input_shape)

    def call(self, inputs):

        sentence_trans = tf.transpose(inputs, [0, 2, 1])
        at = tf.matmul(self.Ws, sentence_trans)
        at = tf.math.tanh(at)
        at = K.exp(at - K.max(at, axis=-1, keepdims=True))
        at = at / K.sum(at, axis=-1, keepdims=True)
        v = K.batch_dot(at, inputs)

        return v

我想在火炬中实现同样的功能；我已经完成了前向传递块，但对如何在 PyTorch 中进行与上述层相同的嵌入和权重初始化感到困惑?

I want to implement the same in the torch; I have already done the forward pass block but am confused about how to do the embedding and weight initialization the same as the above layer in PyTorch?

class Attention_module(torch.nn.Module):
    def __init__(self, class_num):

        # how to initialize weight with same as above keras layer?

    def forward(self, inputs):
        
        sentence_trans = inputs.permute(0, 2, 1)

        at = torch.mm(self.Ws, sentence_trans)
        at = torch.nn.Tanh(at)
        at = torch.exp(at - torch.max(torch.Tensor(at), dim=-1, keepdims=True).values)
        at = at / torch.sum(at, dim = -1, keepdims=True)
        v = torch.einsum('ijk,ikl->ijl', at, inputs)

        return v

谢谢！

推荐答案

class Attention_module(torch.nn.Module):
    def __init__(self, class_num, input_shape):
        super().__init__()
        self.class_num = class_num
        embedding_length = int(input_shape[2])
        self.Ws = torch.nn.Embedding(num_embeddings=class_num, 
                                     embedding_dim=embedding_length)  # Embedding layer
        torch.nn.init.xavier_uniform_(self.Ws.weight) # Glorot initialization

这里是层初始化方法的参考. Xavier init 是另一个名字用于 Glorot 初始化.

Here's the reference for layer initialization methods. Xavier init is another name for Glorot init.

torch.nn.init.xavier_uniform_末尾的_是pytorch约定，表示就地操作.

The _ at the end of torch.nn.init.xavier_uniform_ is a pytorch convention that signifies an inplace operation.

您也可以在运行时使用 torch.nn.init.它不必在 __init__() 内.喜欢:

You can also use torch.nn.init at runtime. It doesn't have to be within __init__(). Like:

att = Attention_module(class_num, input_shape)
torch.nn.init.xavier_uniform_(att.Ws.weight)

或:

for param in att.parameters():
    torch.nn.init.xavier_uniform_(param)

这篇关于从 keras 到 pytorch 的自定义层的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！