搭建神经网络(torch.nn的用法)

零零碎碎总结了一些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 的一般用法:

python 复制代码
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设备并行计算,从而加快模型训练和推断的速度,并提高深度学习任务的效率和扩展性。

相关推荐
*JOKER13 分钟前
Flow Matching&生成算法
人工智能·深度学习·机器学习·大模型·生成模型·flow matching
databook14 分钟前
别让AI代码,变成明天的技术债
人工智能·程序员·ai编程
人工智能AI技术35 分钟前
计算机专业面试必看!90%学生都踩过的算法面雷区
人工智能·面试
千寻girling35 分钟前
面试官 : “ 说一下 Python 中的常用的 字符串和数组 的 方法有哪些 ? ”
人工智能·后端·python
m0_706653231 小时前
深入理解Gumbel-Softmax技巧及其应用
人工智能
七夜zippoe1 小时前
OpenClaw 接入 WhatsApp:消息推送实战
大数据·人工智能·microsoft·whatsapp·openclaw
眠りたいです1 小时前
使用LangChain进行AI应用构建-RAG及相关核心组件认识(二)
人工智能·langchain·rag
WeeJot嵌入式1 小时前
NVIDIA GTC 2026实战:Rubin平台AI五层架构部署指南
人工智能·架构
我材不敲代码1 小时前
OpenCV实战:全自动答题卡识别与评分系统
人工智能·opencv·计算机视觉
SmartBrain1 小时前
AI深度解析:智能体产品核心理念与技术架构
人工智能·架构·aigc