7、深入剖析PyTorch nn.Module源码

🕗 发布于 2024-11-25 17:18 pytorch python numpy 深度学习 机器学习

文章目录

1. 重要类
2. add_modules
3. Apply(fn)
4. register_buffer
5. nn.Parameters®ister_parameters
6. 后续测试

1. 重要类

nn.module --> 所有神经网络的父类，自定义神经网络需要继承此类，并且自定义__init__,forward函数即可：

#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :MyModelNet.py
# @Time      :2024/11/20 13:38
# @Author    :Jason Zhang
import torch
from torch import nn


class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork,self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28 * 28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10)
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits


if __name__ == "__main__":
    run_code = 0
    x_row = 28
    x_column = 28
    x_total = x_row * x_column
    x = torch.arange(x_total, dtype=torch.float).reshape((1, x_row, x_column))
    my_net = NeuralNetwork()
    y = my_net(x)
    print(f"y.shape={y.shape}")
    print(my_net)

结果：

y.shape=torch.Size([1, 10])
NeuralNetwork(
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (linear_relu_stack): Sequential(
    (0): Linear(in_features=784, out_features=512, bias=True)
    (1): ReLU()
    (2): Linear(in_features=512, out_features=512, bias=True)
    (3): ReLU()
    (4): Linear(in_features=512, out_features=10, bias=True)
  )
)

2. add_modules

通过add_modules在旧的网络里面添加新的网络

重点：用nn.ModuleList自带的insert,新的网络继承自老网络中，直接用按位置插入
python

import torch
from torch import nn
from pytorch_model_summary import summary

torch.manual_seed(2323)


class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.flatten = nn.Flatten()
        self.block = nn.ModuleList([
            nn.Linear(28 * 28, 512),
            nn.ReLU(),
            nn.Linear(512, 10)
        ])

    def forward(self, x):
        x = self.flatten(x)
        for layer in self.block:
            x = layer(x)
        return x


class MyNewNet(MyModel):
    def __init__(self):
        super(MyNewNet, self).__init__()
        self.block.insert(2, nn.Linear(512, 256))  # 插入新层
        self.block.insert(3, nn.ReLU())  # 插入新的激活函数
        self.block.insert(4, nn.Linear(256, 512))  # 插入另一层
        self.block.insert(5, nn.ReLU())  # 插入激活函数


if __name__ == "__main__":
    # 测试原始模型
    my_model = MyModel()
    print("Original Model:")
    print(summary(my_model, torch.ones((1, 28, 28))))

    # 测试新模型
    my_new_model = MyNewNet()
    print("\nNew Model:")
    print(summary(my_new_model, torch.ones((1, 28, 28))))

结果：

Original Model:
-----------------------------------------------------------------------
      Layer (type)        Output Shape         Param #     Tr. Param #
=======================================================================
         Flatten-1            [1, 784]               0               0
          Linear-2            [1, 512]         401,920         401,920
            ReLU-3            [1, 512]               0               0
          Linear-4             [1, 10]           5,130           5,130
=======================================================================
Total params: 407,050
Trainable params: 407,050
Non-trainable params: 0
-----------------------------------------------------------------------

New Model:
-----------------------------------------------------------------------
      Layer (type)        Output Shape         Param #     Tr. Param #
=======================================================================
         Flatten-1            [1, 784]               0               0
          Linear-2            [1, 512]         401,920         401,920
            ReLU-3            [1, 512]               0               0
          Linear-4            [1, 256]         131,328         131,328
            ReLU-5            [1, 256]               0               0
          Linear-6            [1, 512]         131,584         131,584
            ReLU-7            [1, 512]               0               0
          Linear-8             [1, 10]           5,130           5,130
=======================================================================
Total params: 669,962
Trainable params: 669,962
Non-trainable params: 0
-----------------------------------------------------------------------

3. Apply(fn)

模型权重weight,bias 的初始化

python

import torch.nn as nn
import torch


class MyAwesomeModel(nn.Module):
    def __init__(self):
        super(MyAwesomeModel, self).__init__()
        self.fc1 = nn.Linear(3, 4)
        self.fc2 = nn.Linear(4, 5)
        self.fc3 = nn.Linear(5, 6)


# 定义初始化函数
@torch.no_grad()
def init_weights(m):
    print(m)
    if type(m) == nn.Linear:
        m.weight.fill_(1.0)
        print(m.weight)


# 创建神经网络实例
model = MyAwesomeModel()

# 应用初始化权值函数到神经网络上
model.apply(init_weights)

结果：

Linear(in_features=3, out_features=4, bias=True)
Parameter containing:
tensor([[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]], requires_grad=True)
Linear(in_features=4, out_features=5, bias=True)
Parameter containing:
tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]], requires_grad=True)
Linear(in_features=5, out_features=6, bias=True)
Parameter containing:
tensor([[1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.]], requires_grad=True)
MyAwesomeModel(
  (fc1): Linear(in_features=3, out_features=4, bias=True)
  (fc2): Linear(in_features=4, out_features=5, bias=True)
  (fc3): Linear(in_features=5, out_features=6, bias=True)
)

Process finished with exit code 0

4. register_buffer

将模型中添加常数项。比如加1

python:

#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :RegisterBuffer.py
# @Time      :2024/11/23 19:21
# @Author    :Jason Zhang
import torch
from torch import nn


class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.register_buffer("my_buffer_a", torch.ones(2, 3))

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


if __name__ == "__main__":
    run_code = 0
    my_test = MyNet()
    in_x = torch.arange(6).reshape((2, 3))
    y = my_test(in_x)
    print(f"x=\n{in_x}")
    print(f"y=\n{y}")

结果：

x=
tensor([[0, 1, 2],
        [3, 4, 5]])
y=
tensor([[1., 2., 3.],
        [4., 5., 6.]])

5. nn.Parameters&register_parameters

python

#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :ParameterTest.py
# @Time      :2024/11/23 19:37
# @Author    :Jason Zhang
import torch
from torch import nn


class MyModule(nn.Module):
    def __init__(self, in_size, out_size):
        self.in_size = in_size
        self.out_size = out_size
        super(MyModule, self).__init__()
        self.test = torch.rand(self.in_size, self.out_size)
        self.linear = nn.Linear(self.in_size, self.out_size)

    def forward(self, x):
        x = self.linear(x)
        return x


class MyModuleRegister(nn.Module):
    def __init__(self, in_size, out_size):
        self.in_size = in_size
        self.out_size = out_size
        super(MyModuleRegister, self).__init__()
        self.test = torch.rand(self.in_size, self.out_size)
        self.linear = nn.Linear(self.in_size, self.out_size)

    def forward(self, x):
        x = self.linear(x)
        return x


class MyModulePara(nn.Module):
    def __init__(self, in_size, out_size):
        self.in_size = in_size
        self.out_size = out_size
        super(MyModulePara, self).__init__()
        self.test = nn.Parameter(torch.rand(self.in_size, self.out_size))
        self.linear = nn.Linear(self.in_size, self.out_size)

    def forward(self, x):
        x = self.linear(x)
        return x


if __name__ == "__main__":
    run_code = 0
    test_in = 4
    test_out = 6
    my_test = MyModule(test_in, test_out)
    my_test_para = MyModulePara(test_in, test_out)
    test_list = list(my_test.named_parameters())
    test_list_para = list(my_test_para.named_parameters())
    my_test_register = MyModuleRegister(test_in, test_out)
    para_register = nn.Parameter(torch.rand(test_in, test_out))
    my_test_register.register_parameter('para_add_register', para_register)
    test_list_para_register = list(my_test_register.named_parameters())

    print(f"*" * 50)
    print(f"test_list=\n{test_list}")
    print(f"*" * 50)
    print(f"*" * 50)
    print(f"test_list_para=\n{test_list_para}")
    print(f"*" * 50)
    print(f"*" * 50)
    print(f"test_list_para_register=\n{test_list_para_register}")
    print(f"*" * 50)

结果：

**************************************************
test_list=
[('linear.weight', Parameter containing:
tensor([[ 0.3805, -0.3368,  0.2348,  0.4525],
        [-0.4557, -0.3344,  0.1368, -0.3471],
        [-0.3961,  0.3302,  0.1904, -0.0111],
        [ 0.4542, -0.3325, -0.3782,  0.0376],
        [ 0.2083, -0.3113, -0.3447, -0.1503],
        [ 0.0343,  0.0410, -0.4216, -0.4793]], requires_grad=True)), ('linear.bias', Parameter containing:
tensor([-0.3465, -0.4510,  0.4919,  0.1967, -0.1366, -0.2496],
       requires_grad=True))]
**************************************************
**************************************************
test_list_para=
[('test', Parameter containing:
tensor([[0.1353, 0.9934, 0.0462, 0.2103, 0.3410, 0.0814],
        [0.7509, 0.2573, 0.8030, 0.0952, 0.1381, 0.5360],
        [0.1972, 0.1241, 0.5597, 0.2691, 0.3226, 0.0660],
        [0.3333, 0.8031, 0.9226, 0.4290, 0.3660, 0.6159]], requires_grad=True)), ('linear.weight', Parameter containing:
tensor([[-0.0633, -0.4030, -0.4962,  0.1928],
        [-0.1707,  0.2259,  0.0373, -0.0317],
        [ 0.4523,  0.2439, -0.1376, -0.3323],
        [ 0.3215,  0.1283,  0.0729,  0.3912],
        [ 0.0262, -0.1087,  0.4721, -0.1661],
        [-0.1055, -0.2199, -0.4974, -0.3444]], requires_grad=True)), ('linear.bias', Parameter containing:
tensor([ 0.3702, -0.0142, -0.2098, -0.0910, -0.2323, -0.0546],
       requires_grad=True))]
**************************************************
**************************************************
test_list_para_register=
[('para_add_register', Parameter containing:
tensor([[0.2428, 0.1388, 0.6612, 0.4215, 0.0215, 0.2618],
        [0.4234, 0.0160, 0.8947, 0.4784, 0.4403, 0.4800],
        [0.8845, 0.1469, 0.6894, 0.7050, 0.5911, 0.7702],
        [0.7694, 0.0491, 0.3583, 0.4451, 0.2282, 0.4293]], requires_grad=True)), ('linear.weight', Parameter containing:
tensor([[ 0.1358, -0.4704, -0.4181, -0.4504],
        [ 0.0903,  0.3235, -0.3164, -0.4163],
        [ 0.1342,  0.3108,  0.0612, -0.2910],
        [ 0.3527,  0.3397, -0.0414, -0.0408],
        [-0.4877,  0.1925, -0.2912, -0.2239],
        [-0.0081, -0.1730,  0.0921, -0.4210]], requires_grad=True)), ('linear.bias', Parameter containing:
tensor([-0.2194,  0.2233, -0.4950, -0.3260, -0.0206, -0.0197],
       requires_grad=True))]
**************************************************

6. 后续测试

register_module
get_submodule
get_parameter

原文地址：https://blog.csdn.net/scar2016/article/details/143990867

免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！

上一篇：【深入理解RabbitMQ】七大工作模式
下一篇：数据集-目标检测系列- 蘑菇检测数据集 mushroom ＞＞ DataBall

移动充储机器人“小奥”的多场景应用（上）
此外，“小奥”在设计上还考虑了环境适应性，它能够在各种天气条件下稳定工作，无论是炎热的夏季还是寒冷的冬季，都能保证充电服务的连续性和可靠性。它采用的充电技术减少了对电网的冲击，优化了能源的使用效率，有
阅读更多2024-11-26
day07（单片机高级）继电器模块绘制
继电器模块绘制
阅读更多2024-11-26
react函数式组件中的路由传参方式
useParamsstateNavigate路由传参。React Router 提供了多种方式来传递和获取路由参数。
阅读更多2024-11-26
MySQL中的锁与优化SQL查询性能
MySQL作为一种高效、稳定、易用的开源关系型数据库管理系统（RDBMS），在大数据量和高并发的场景中，其性能优化显得尤为重要。锁机制和SQL查询优化是MySQL性能调优的两个关键方面。本文将详细探讨
阅读更多2024-11-26
Cesium教程06_材质设置
cesium设置材质
阅读更多2024-11-26
Three.js CSS2D/CSS3D渲染器
CSS2DRenderer 渲染器用于在 3D 场景中渲染纯 2D 的 HTML 元素。这些元素不会具有 3D 透视效果，但可以与 3D 对象一起移动和旋转，非常适合用于标签、注释等需要固定显示的内容
阅读更多2024-11-26
【过滤器】一文了解 .NET Core 中各种 Filter
ASP.NET Core中，Filter（过滤器）是一种强大的机制，允许你在请求处理管道中的特定阶段运行代码。这些过滤器提供了在请求的不同阶段执行逻辑的能力，比如授权、日志记录、异常处理等。ASP.N
阅读更多2024-11-26
人工智能-深度学习-Torch框架-手动构建回归流程
导入函数，用于生成回归数据集。: 导入math模块，用于进行数学计算，例如向上取整。: 导入random模块，用于随机打乱数据集。: 导入torch库，用于张量操作和神经网络训练。
阅读更多2024-11-26
android 11添加切换分屏功能
自Android 7开始官方就支持分屏显示，但没有切换分屏的功能，即交换上下屏幕。直到Android 13开始才支持切换分屏，操作方式是：分屏模式下双击中间分割线就会交换上下屏位置。本文的目的就是在A
阅读更多2024-11-26
高防服务器HOT：网络安全的无形盾牌，护航业务稳定
高防服务器通过卓越的防护能力，为企业筑起一道坚实的安全屏障，不仅能保障业务稳定，更能为企业争取宝贵的声誉和客户信赖。高防服务器不仅能够提供强大的抗攻击能力，还兼具高性能网络和低延迟特性，广泛应用于对网
阅读更多2024-11-26