码迷,mamicode.com
首页 > 其他好文 > 详细

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

时间:2020-07-17 13:40:24      阅读:103      评论:0      收藏:0      [点我收藏+]

标签:gpu   size   numpy   自动   tensor   dev   ble   env   设置   

先设置环境

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()

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

标签:gpu   size   numpy   自动   tensor   dev   ble   env   设置   

原文地址:https://www.cnblogs.com/baitian963/p/13329308.html

(0)
(0)
   
举报
评论 一句话评论(0
登录后才能评论!
© 2014 mamicode.com 版权所有  联系我们:gaon5@hotmail.com
迷上了代码!