nvlink 训练笔记

目录

还没测试出效果


还没测试出效果

python 复制代码
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR10
from torchvision.transforms import ToTensor

# 定义上述的大型全连接层模型
class LargeFullyConnectedModel(nn.Module):
    def __init__(self):
        super(LargeFullyConnectedModel, self).__init__()
        input_size = 10000
        hidden_size1 = 20000
        hidden_size2 = 15000
        hidden_size3 = 12000
        output_size = 5000

        self.fc1 = nn.Linear(input_size, hidden_size1)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size1, hidden_size2)
        self.relu2 = nn.ReLU()
        self.fc3 = nn.Linear(hidden_size2, hidden_size3)
        self.relu3 = nn.ReLU()
        self.fc4 = nn.Linear(hidden_size3, output_size)

    def forward(self, x):
        x = self.relu1(self.fc1(x))
        x = self.relu2(self.fc2(x))
        x = self.relu3(self.fc3(x))
        x = self.fc4(x)
        return x

# 初始化模型并准备多卡环境
devices = [0, 1]  # 指定要使用的显卡编号列表
model = LargeFullyConnectedModel()
if torch.cuda.device_count() > 1 and len(devices) > 1:
    print(f"使用 {len(devices)} 个 GPU 进行推理")
    model = nn.DataParallel(model, device_ids=devices)
else:
    print("仅使用单个 GPU 进行推理")
model.to(torch.device(f"cuda:{devices[0]}" if torch.cuda.is_available() else "cpu"))

# 模拟数据加载(这里只是示例,实际需根据你的数据进行调整)
batch_size = 32
input_size = 10000
data = torch.randn(batch_size, input_size).to(torch.device(f"cuda:{devices[0]}"))
targets = torch.randint(0, 5000, (batch_size,)).to(torch.device(f"cuda:{devices[0]}"))

# 定义推理函数
def inference():
    model.eval()
    with torch.no_grad():
        outputs = model(data)
        # 可以根据需要进行后续处理,如计算损失、准确率等
    return outputs

if __name__ == "__main__":
    inference()
相关推荐
辛勤的程序猿27 分钟前
YOLO即插即用---PConv
深度学习·yolo·计算机视觉
-喵侠客-1 小时前
探索开源MiniMind项目:让大语言模型不再神秘(1)
人工智能·深度学习·语言模型·自然语言处理
新手小白勇闯新世界1 小时前
论文阅读-用于图像识别的深度残差学习
论文阅读·人工智能·深度学习·学习·计算机视觉
新手小白勇闯新世界1 小时前
论文阅读- --DeepI2P:通过深度分类进行图像到点云配准
论文阅读·深度学习·算法·计算机视觉
子午1 小时前
【车辆车型识别】Python+卷积神经网络算法+深度学习+人工智能+TensorFlow+算法模型
人工智能·python·深度学习
chan_lay1 小时前
图论导引 - 目录、引言、第一章 - 11/05
笔记·图论
B站计算机毕业设计超人2 小时前
计算机毕业设计Hadoop+大模型地震预测系统 地震数据分析可视化 地震爬虫 大数据毕业设计 Spark 机器学习 深度学习 Flink 大数据
大数据·hadoop·爬虫·深度学习·机器学习·数据分析·课程设计
、你的风衣2 小时前
FITS论文解析
人工智能·深度学习·论文笔记
亦枫Leonlew2 小时前
三维测量与建模笔记 - 3.1 相机标定基本概念
笔记·三维重建·1024程序员节·射影几何