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()
相关推荐
九章云极AladdinEdu4 小时前
超参数自动化调优指南:Optuna vs. Ray Tune 对比评测
运维·人工智能·深度学习·ai·自动化·gpu算力
Hello_Embed4 小时前
STM32HAL 快速入门(二十):UART 中断改进 —— 环形缓冲区解决数据丢失
笔记·stm32·单片机·学习·嵌入式软件
咸甜适中5 小时前
rust语言 (1.88) 学习笔记:客户端和服务器端同在一个项目中
笔记·学习·rust
Grassto5 小时前
RAG 从入门到放弃?丐版 demo 实战笔记(go+python)
笔记
Magnetic_h6 小时前
【iOS】设计模式复习
笔记·学习·ios·设计模式·objective-c·cocoa
研梦非凡6 小时前
ICCV 2025|从粗到细:用于高效3D高斯溅射的可学习离散小波变换
人工智能·深度学习·学习·3d
通街市密人有8 小时前
IDF: Iterative Dynamic Filtering Networks for Generalizable Image Denoising
人工智能·深度学习·计算机视觉
智数研析社9 小时前
9120 部 TMDb 高分电影数据集 | 7 列全维度指标 (评分 / 热度 / 剧情)+API 权威源 | 电影趋势分析 / 推荐系统 / NLP 建模用
大数据·人工智能·python·深度学习·数据分析·数据集·数据清洗
周周记笔记9 小时前
学习笔记:第一个Python程序
笔记·学习
丑小鸭是白天鹅9 小时前
Kotlin协程详细笔记之切线程和挂起函数
开发语言·笔记·kotlin