PyTorch的nn.Module类在TensorFlow框架中的等效替代是什么？

PyTorch 中用于自定义网络模块的 nn.Module，在 TensorFlow 中对应的基类是 tf.keras.layers.Layer，你可以通过继承该基类实现自定义层/模块。

你给出的 ChannelPool 类对应的 TensorFlow 实现如下：

import tensorflow as tf

class ChannelPool(tf.keras.layers.Layer):
    def call(self, x):
        # 以下逻辑对齐你给出的PyTorch代码，适配通道优先格式（输入维度为 [batch, channels, height, width]）
        max_pool = tf.reduce_max(x, axis=1, keepdims=True)
        mean_pool = tf.reduce_mean(x, axis=1, keepdims=True)
        return tf.concat([max_pool, mean_pool], axis=1)

补充说明：

• 继承 tf.keras.layers.Layer 后，PyTorch 中定义前向逻辑的 forward 方法对应替换为 call 方法即可。
• 如果你使用 TensorFlow 默认的通道最后格式（输入维度为 [batch, height, width, channels]），只需要把上述代码中所有 axis=1 改为 axis=-1 即可。
• 自定义的 ChannelPool 层可以和 TensorFlow 内置层完全一样的方式，用在 Sequential、函数式API构建的模型中，和 PyTorch 中继承 nn.Module 的模块用法对齐。

内容的提问来源于stack exchange，提问作者Mithil Neeraj Salunkhe