把张量计算放到GPU上进行计算

Posted 2020-12-16 baitian963

先设置环境

parser = argparse.ArgumentParser()
parser.add_argument(‘--gpu‘,type=str,default=‘7‘)
opt=parser.parse_args()
print(opt)
os.environ["CUDA_VISIBLE_DEVICES"]=opt.gpu



然后把要放到GPU中的变量变成cuda变量

image_torch = image_torch.cuda()
其中变量要为tensor才行

计算过程中自然产生的变量不用再.cuda（）放到GPU上，已经自动放上去了

但是新声明的变量要放上去，比如中间过程中：

S_out = torch.zeros(S_max.size())
就需要S_out = S_out.cuda（）放到GPU上



最后把计算结果放回到CPU上

S_mean = S_mean.cpu()
S_mean = S_mean.numpy().tolist()