Блог

Я пытаюсь реализовать триплетное внимание в TensorFlow. Один из вопросов, с которым я сталкиваюсь, заключается в том, что использовать вместо NN.module в TensorFlow

 class ChannelPool(nn.Module):
    def forward(self, x):
        return torch.cat( (torch.max(x,1)[0].unsqueeze(1), torch.mean(x,1).unsqueeze(1)), dim=1)
 

Что я должен поставить nn.Module здесь на место ?

В этом случае nn.Module используется для создания пользовательского слоя. У TensorFlow есть учебник по этому вопросу, пожалуйста, взгляните. Короче говоря, один из способов , которым вы могли бы это реализовать, — это tf.keras.layers.Layer , где call эквивалент forward в PyTorch:

 class ChannelPool(tf.keras.layers.Layer):
    def call(self, inputs):
        return tf.concat((tf.reduce_max(inputs, axis=1, keepdims=True), tf.reduce_mean(inputs, axis=1, keepdims=True)), axis=1)
 

Вы можете проверить, что они эквивалентны, как это:

 import torch
from torch import nn
import tensorflow as tf
import numpy as np

class PyTorch_ChannelPool(nn.Module):
    def forward(self, x):
        return torch.cat((torch.max(x, 1)[0].unsqueeze(1), torch.mean(x, 1).unsqueeze(1)), dim=1)

class TensorFlow_ChannelPool(tf.keras.layers.Layer):
    def call(self, inputs):
        return tf.concat((tf.reduce_max(inputs, axis=1, keepdims=True), tf.reduce_mean(inputs, axis=1, keepdims=True)), axis=1)

np.random.seed(2021)
x = np.random.random((1,2,3,4)).astype(np.float32)

a = PyTorch_ChannelPool()
b = TensorFlow_ChannelPool()

pytorch_output = a(torch.from_numpy(x)).numpy()
tensorflow_output = b(x).numpy()

np.all(pytorch_output == tensorflow_output)
# >>> True