PyTorch 与 TensorFlow 中基于自定义层的 DNN 实现对比

深度学习双雄对决：PyTorch vs TensorFlow 自定义层大比拼

一、TensorFlow 实现 DNN

1. 核心逻辑

• 直接继承 tf.keras.layers.Layer：无需中间类，直接在 build 中定义多层结构。
• 动态参数管理：通过 add_weight 注册每一层的权重和偏置。

import tensorflow as tf

class CustomDNNLayer(tf.keras.layers.Layer):
    def __init__(self, hidden_units, output_dim, **kwargs):
        super(CustomDNNLayer, self).__init__(**kwargs)
        self.hidden_units = hidden_units
        self.output_dim = output_dim

    def build(self, input_shape):
        # 输入层到第一个隐藏层
        self.w1 = self.add_weight(
            name='w1', 
            shape=(input_shape[-1], self.hidden_units[0]),
            initializer='random_normal',
            trainable=True
        )
        self.b1 = self.add_weight(
            name='b1',
            shape=(self.hidden_units[0],),
            initializer='zeros',
            trainable=True
        )

        # 隐藏层之间
        self.ws = []
        self.bs = []
        for i in range(len(self.hidden_units) - 1):
            self.ws.append(self.add_weight(
                name=f'w{i+2}', 
                shape=(self.hidden_units[i], self.hidden_units[i+1]),
                initializer='random_normal',
                trainable=True
            ))
            self.bs.append(self.add_weight(
                name=f'b{i+2}',
                shape=(self.hidden_units[i+1],),
                initializer='zeros',
                trainable=True
            ))

        # 输出层
        self.wo = self.add_weight(
            name='wo',
            shape=(self.hidden_units[-1], self.output_dim),
            initializer='random_normal',
            trainable=True
        )
        self.bo = self.add_weight(
            name='bo',
            shape=(self.output_dim,),
            initializer='zeros',
            trainable=True
        )

    def call(self, inputs):
        x = tf.matmul(inputs, self.w1) + self.b1
        x = tf.nn.relu(x)

        for i in range(len(self.hidden_units) - 1):
            x = tf.matmul(x, self.ws[i]) + self.bs[i]
            x = tf.nn.relu(x)

        x = tf.matmul(x, self.wo) + self.bo
        return x

二、PyTorch 实现自定义层

1. 核心逻辑

• 继承 nn.Module：自定义层本质是模块的组合。
• 使用 nn.ModuleList：动态管理多个 nn.Linear 层。

import torch
import torch.nn as nn

class CustomPyTorchDNN(nn.Module):
    def __init__(self, input_size, hidden_sizes, output_size):
        super(CustomPyTorchDNN, self).__init__()
        self.hidden_layers = nn.ModuleList()
        prev_size = input_size

        # 动态添加隐藏层
        for hidden_size in hidden_sizes:
            self.hidden_layers.append(nn.Linear(prev_size, hidden_size))
            prev_size = hidden_size

        # 输出层
        self.output_layer = nn.Linear(prev_size, output_size)

    def forward(self, x):
        for layer in self.hidden_layers:
            x = torch.relu(layer(x))
        x = self.output_layer(x)
        return x

三、关键差异对比

四、总结

• TensorFlow：通过直接继承 tf.keras.layers.Layer，可实现完全自定义的 DNN，但需手动管理多层权重和激活逻辑，适合对模型细节有严格控制需求的场景。
• PyTorch：通过直接继承 nn.Module，可实现完全自定义的 DNN;利用 nn.ModuleList 和 nn.Linear 的组合，能高效构建标准 DNN 结构，代码简洁且易于扩展，适合快速原型开发和研究场景。

两种实现均满足用户对“直接继承核心类 + 使用基础组件”的要求，可根据具体任务选择框架。