pytorch使用技巧

最新推荐文章于 2024-07-09 15:30:57 发布

wxzjfdcx

最新推荐文章于 2024-07-09 15:30:57 发布

阅读量890

点赞数 2

分类专栏：深度学习

本文链接：https://blog.csdn.net/wxzjfdcx/article/details/101630216

版权

指定GPU

推荐使用方法二

方法一：

使用torch的函数

torch.cuda.set_device(gpu_id)

通过gpu_id指定使用的GPU编号

方法二：

使用设置环境变量的方法
在python代码中：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "gpu_id"

或者在linux中直接设置环境变量

export CUDA_VISIBLE_DEVICES=gpu_id

数据并行nn.DataParallel

class Net(nn.Module):
	def __init__(self):
		super().__init__()
		self.dense = nn.Linear(5, 5)
	def forward(self, data):
		print('forward data shape:' , data.shape)
		out = self.dense(data)
		return out

model = Net()
model = nn.DataParallel(model.cuda(), device_ids=[0,1])

data = torch.rand(10, 5).cuda()
out = model(data)
print('out data shape: ', out.shape)