GPU程序cuda设置

程序中cuda的设置

默认会用 Cuda 0，有时候容易拥堵

device = "cuda:0" if torch.cuda.is_available() else "cpu"
print('device = ', device)
model = net().to(device)

修改方法

一般使用.cuda()或者.to(device)，注意device的名字不能直接用数字

而是用torch.device(‘cuda:1’)表示

device = torch.device("cuda:0")
device1 = torch.device("cuda:1") 
for batch_idx, (img, label) in enumerate(train_loader):
 img=img.to(device1)
 label=label.to(device1)

实例化网络

model = Net()
model.to(device) #使用序号为0的GPU
# 或model.to(device1) #使用序号为1的GPU

把tensor改成numpy

measurements, channels, sigma_2 = load_CEdataset(measurement_path, channel_path)
measurements = measurements.to(device)
channels = channels.to(device)
sigma_2 = sigma_2.to(device)

measurements = measurements[0:500,:]
measurements = measurements.cpu().numpy()
channels = channels[0:500,:]
channels = channels.cpu().numpy()
sigma_2 = sigma_2[0:500,:]
sigma_2 = sigma_2.cpu().numpy()

多GPU

假设有4个GPU,其id设置如下

device_ids =[0,1,2,3]
#对数据
input_data=input_data.to(device=device_ids[0])
#对于模型
net = torch.nn.DataParallel(model)
net.to(device)

此时利用某一些gpu来存放数据（分配来做一些事），DataParallel会自动让所有gpu进行参与