PyTorch的forward函数在TensorFlow中的对应实现是什么？

在PyTorch中，forward方法是nn.Module子类的核心，定义了模型/层的前向传播逻辑。在TensorFlow中，对应的实现方式分以下几种场景：

1. 自定义模型类（对应PyTorch的nn.Module） #

TensorFlow中继承tf.keras.Model类，通过重写call方法来实现前向传播逻辑，这和PyTorch的forward完全等价。

PyTorch示例 #

import torch
import torch.nn as nn

class PyTorchCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 16, kernel_size=3)
        self.relu = nn.ReLU()
        self.flatten = nn.Flatten()
        self.fc = nn.Linear(16 * 30 * 30, 10)  # 假设输入为32x32 RGB图像

    def forward(self, x):
        x = self.conv(x)
        x = self.relu(x)
        x = self.flatten(x)
        x = self.fc(x)
        return x

对应TensorFlow实现 #

import tensorflow as tf

class TensorFlowCNN(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.conv = tf.keras.layers.Conv2D(16, kernel_size=3, input_shape=(32, 32, 3))
        self.relu = tf.keras.layers.ReLU()
        self.flatten = tf.keras.layers.Flatten()
        self.fc = tf.keras.layers.Dense(10)

    def call(self, x):
        x = self.conv(x)
        x = self.relu(x)
        x = self.flatten(x)
        x = self.fc(x)
        return x

使用方式和PyTorch一致，直接实例化模型后传入输入数据：

# TensorFlow调用示例
model = TensorFlowCNN()
output = model(tf.random.normal((1, 32, 32, 3)))

2. 自定义层类（对应PyTorch的自定义nn.Module层） #

如果是自定义可训练层，TensorFlow中继承tf.keras.layers.Layer，同样重写call方法实现前向逻辑。和PyTorch不同的是，TensorFlow通常在build方法中根据输入形状初始化权重（第一次接收输入时自动触发）。

PyTorch自定义层示例 #

class PyTorchLinear(nn.Module):
    def __init__(self, out_features):
        super().__init__()
        self.out_features = out_features

    def forward(self, x):
        # 懒初始化权重（PyTorch 1.10+支持）
        if not hasattr(self, 'weight'):
            self.weight = nn.Parameter(torch.randn(x.shape[-1], self.out_features))
            self.bias = nn.Parameter(torch.randn(self.out_features))
        return x @ self.weight + self.bias

对应TensorFlow自定义层 #

class TensorFlowLinear(tf.keras.layers.Layer):
    def __init__(self, out_features):
        super().__init__()
        self.out_features = out_features

    def build(self, input_shape):
        # 根据输入形状初始化可训练参数
        self.weight = self.add_weight(
            shape=(input_shape[-1], self.out_features),
            initializer='random_normal',
            trainable=True
        )
        self.bias = self.add_weight(
            shape=(self.out_features,),
            initializer='random_normal',
            trainable=True
        )

    def call(self, x):
        return tf.matmul(x, self.weight) + self.bias

3. 函数式API（TensorFlow特有方式） #

如果不需要定义类，TensorFlow的函数式API可以直接通过组合层来构建前向传播逻辑，无需显式写call或forward方法。

def build_cnn_model():
    inputs = tf.keras.Input(shape=(32, 32, 3))
    x = tf.keras.layers.Conv2D(16, kernel_size=3)(inputs)
    x = tf.keras.layers.ReLU()(x)
    x = tf.keras.layers.Flatten()(x)
    outputs = tf.keras.layers.Dense(10)(x)
    return tf.keras.Model(inputs=inputs, outputs=outputs)

# 使用示例
model = build_cnn_model()
output = model(tf.random.normal((1, 32, 32, 3)))

关键对应关系总结 #

内容的提问来源于stack exchange，提问作者Tim Bretschneider