需要在Pytorch下跑一SalGAN,但是代码居然用CPU来计算,将其改为GPU版。Pytorch下使用GPU计算需要三步,记录一下:
1.将数据载入GPU:
input -> input.cuda()
2.将模型载入GPU
model = model.cuda()
3.将输出的预测转为numpy格式
result.data.numpy() -> result.data.cpu().numpy()
需要在Pytorch下跑一SalGAN,但是代码居然用CPU来计算,将其改为GPU版。Pytorch下使用GPU计算需要三步,记录一下:
1.将数据载入GPU:
input -> input.cuda()
2.将模型载入GPU
model = model.cuda()
3.将输出的预测转为numpy格式
result.data.numpy() -> result.data.cpu().numpy()