pytorch的python API略读--torch.nn.Module

机器视觉全栈er

于 2022-04-24 20:18:19 发布

阅读量1.2k

点赞数

分类专栏： pytorch的python API 文章标签： pytorch python 深度学习计算机视觉人工智能

本文链接：https://blog.youkuaiyun.com/qq_42322191/article/details/124390842

版权

本文介绍了PyTorch中的torch.nn.Module类，它是所有网络模块的基类，用于构建复杂的深度学习模型。Module类支持模块嵌套，提供eval()和train()方法来切换训练和评估模式，并且需要重写forward方法来定义前向传播过程。此外，还提到了Module类的Hook函数，如register_forward_hook等，它们可以在特定时刻执行自定义操作。

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

作者：机器视觉全栈er
网站：www.cvtutorials.com

2.3 torch.nn

本节我们看下torch.nn，这个模块是基础模块。

2.3.1 容器

torch.nn.Module：这个类是所有网络模块的基类，甚至你写普通的数学函数的时候，也可以用这个作为基类。内置的模块和你自定义的模块都可以相互嵌套，构成一个树状的模块结构。下面我们来看下这个Module类常用案例：

import torch
import torch.nn as nn
import torch.nn.functional as F

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 3, 3)
        self.conv2 = nn.Conv2d(3, 3