pytorch Module中的forward使用for循环与不使用for循环的区别

最新推荐文章于 2024-02-06 09:14:31 发布

原创最新推荐文章于 2024-02-06 09:14:31 发布 · 1.8k 阅读

2 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #深度学习 #python

本文通过实例探讨了PyTorch中Module的__init__()和forward()函数的区别。作者发现，__init__()定义了Module的网络结构，而forward()决定了网络的执行流程。在forward()的for循环里，输入数据是通过同一个网络进行处理，而非每次都使用新的网络。这有助于初学者理解PyTorch模块的内部工作原理。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

作为初学者，看代码一度很迷惑，module中的forward函数中for循环，输入的Tensor数据是在同一个网络循环，还是依次向前推进了多个不同的网络。于是，我经过了下面的测试。

import torch.nn as nn


class Model1(nn.Module):
    def __init__(self, hidden_dim):
        super().__init__()

        self.block = nn.Linear(hidden_dim, hidden_dim)

    def forward(self, x):
        x = self.block(x) + x
        x = self.block(x) + x
        x = self.block(x) + x
        return x


class Model2(nn.Module):
    def __init__(self, hidden_dim):
        super().__init__()

        self.block = nn.Linear(hidden_dim, hidden_dim)

    def forward(self, x):
        for i in range(3):
            x = self.block(x) + x
        return x


model1 = Model1(hidden_dim=10)
model2=Model2(hidden_dim=10)
print(model1)
print(model2)

得到以下结果：

Model1(
    (block): Linear(in_feature=10, out_feature=10, bias=True)
)
Model2(
    (block): Linear(in_feature=10, out_feature=10, bias=True)
)

然后我就悟了！一个Module的结构到底是由什么构成的，是__init__()还是forward()？结论是__init__()决定了Module有哪些网络，forward()决定了Module的网络是如何连接的。在forward()中无论如何调用__init__()中定义的某个网络，始终都是同一个网络。

那么文章开头那个问题的答案就有了，答案是：for循环中，通过的是同一个网络。