当前位置：首页 > news >正文

搭建神经网络（torch.nn的用法）

news 来源：原创 2025/6/4 14:02:40

零零碎碎总结了一些torch框架里面nn模块的用法，尤其是关于搭建神经网络的

nn.ModuleList

nn.Module

nn.Sequential

nn.Linear

nn.Dropout

nn.Embedding

nn.DataParallel()

将模型封装起来，便于在多个gpu上并行计算，训练或者推理

nn.DataParallel 是一个用于并行计算的 PyTorch 模型包装器。它可以将模型复制到多个GPU设备上，并自动将输入数据划分为多个子批次（mini-batches），分配给不同的GPU进行计算，并且能够有效地将梯度聚合回主模型。

nn.DataParallel 是一个用于在多个GPU上并行执行模型训练或推断的PyTorch模型包装器。下面是 nn.DataParallel 的一般用法：

1.导入必要的库：

import torch
import torch.nn as nn


2.定义模型：

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, 3)
        # 定义其他层和操作

    def forward(self, x):
        # 定义前向传播逻辑
        return out


3.创建模型对象：

model = MyModel()


4.指定要使用的GPU设备：

device_ids = [0, 1, 2]  # 指定要使用的GPU设备的索引


5.使用 nn.DataParallel 包装模型对象：

model = nn.DataParallel(model, device_ids=device_ids)

以上代码将模型对象 model 包装在 nn.DataParallel 中，使用 device_ids 指定要使用的GPU设备的索引列表。
现在，model 将被自动复制到指定的GPU设备上，并且计算会在每个设备上同时进行。输入数据会被自动划分为多个子批次，并分配给不同的GPU设备。在计算梯度时，梯度将被聚合回主模型，并进行参数更新。
请注意，使用 nn.DataParallel 时，一些操作（如 model.parameters() 和 model.to(device)）可能需要进行适当的调整。具体调整取决于您的代码结构和需求。
使用 nn.DataParallel 可以方便地利用多个GPU设备并行计算，从而加快模型训练和推断的速度，并提高深度学习任务的效率和扩展性。

nn.Parameter

在PyTorch中，nn.Parameter 是一个类，它用于将一个张量包装成一个可训练的参数。它是torch.nn.Parameter类的实例。
通常，当你在定义一个PyTorch模型的时候，你会有一些需要在训练过程中更新的参数。使用nn.Parameter可以将这些参数明确地标记为模型的参数，以便优化器在训练时更新它们的值。
nn.Parameter 的常见用法如下：

import torch
import torch.nn as nn

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

        # 定义一个可训练的参数
        self.weight = nn.Parameter(torch.randn(3, 4))

    def forward(self, x):
        # 在前向传播过程中使用参数
        output = torch.matmul(x, self.weight)
        return output

在上面的例子中，self.weight 被包装成了一个 nn.Parameter，这意味着它是模型的一部分，并且在训练时会被优化器更新。你可以像使用普通张量一样使用 self.weight，但它有一些额外的属性，如 .grad 用于存储梯度，.requires_grad 用于指示是否需要计算梯度等。
注意，nn.Parameter 实际上是 torch.nn.Parameter 的一个别名，所以可以直接使用 nn.Parameter，也可以使用 torch.nn.Parameter。例如：
param = nn.Parameter(torch.randn(3, 4))

这个参数可以像普通张量一样使用，同时它也是模型的一部分，会在反向传播中计算梯度并更新。