全部栏目

首页 操作系统 后端开发 数据库 前端开发 移动端开发
您的当前位置：首页正文

PyTorch中指定GPU块号的方法

2024-11-27 来源：个人技术集锦

1、 改变系统变量使得仅目标显卡可见

2、使用torch.cuda接口

在生成网络对象之前执行：torch.cuda.set_device(0)

3、使用pytorch的并行GPU接口

net = torch.nn.DataParallel(model, device_ids=[0,1,2...])

4、初始化模型时

net = Net.cuda(0)
显示全文