标签: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 size numpy 自动 tensor dev ble env 设置
原文地址:https://www.cnblogs.com/baitian963/p/13329308.html