TensorFlow 中 pytorch NN.module 的别名是什么？

Question

我正在尝试在 TensorFlow 中实现 Triplet Attention。我面临的问题之一是在NN.moduleTensorFlow 中使用什么来代替

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？

Answer 1

在本例中，nn.Module用于创建自定义层。TensorFlow 有这方面的教程，请看一下。简而言之，实现它的一种方法是使用tf.keras.layers.Layer，其中call相当于forwardPyTorch 中的：

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