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

使用argparse模块 进行device-agnostic设备未知时的编码

时间:2020-06-26 20:08:04      阅读:72      评论:0      收藏:0      [点我收藏+]

标签:argparse   add   imp   cuda   通过   pre   tor   none   使用   

import argparse
import torch

#核心就是argparse类生成参数包parser实例
parser = argparse.ArgumentParser(description=‘PyTorch Example‘)

#通过这个类实例的方法如add_argument添加参数属性,parse_args解析成点运算可访问的参数包形式
parser.add_argument(‘--disable-cuda‘, action=‘store_true‘,
                    help=‘Disable CUDA‘)     #添加属性disable-cuda

args = parser.parse_args()#参数解析,得到一个类似字典的args

args.device = None

if not args.disable_cuda and torch.cuda.is_available():#参数设置要用cuda,且有GPU可用
    args.device = torch.device(‘cuda‘)
else:
    args.device = torch.device(‘cpu‘)

x = torch.empty((8, 42), device=args.device)
net = Network().to(device=args.device)

使用argparse模块 进行device-agnostic设备未知时的编码

标签:argparse   add   imp   cuda   通过   pre   tor   none   使用   

原文地址:https://www.cnblogs.com/Henry-ZHAO/p/13195759.html

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