本文整理汇总了Python中nn.DataParallel方法的典型用法代码示例。如果您正苦于以下问题：Python nn.DataParallel方法的具体用法？Python nn.DataParallel怎么用？Python nn.DataParallel使用的例子？那么恭喜您, 这里精选的方法代码示例或许可以为您提供帮助。您也可以进一步了解该方法所在类nn的用法示例。
在下文中一共展示了nn.DataParallel方法的5个代码示例，这些例子默认根据受欢迎程度排序。您可以为喜欢或者感觉有用的代码点赞，您的评价将有助于系统推荐出更棒的Python代码示例。
示例1: save_ckpt
# 需要导入模块: import nn [as 别名]
# 或者: from nn import DataParallel [as 别名]
def save_ckpt(output_dir, args, model, optimizer):
"""Save checkpoint"""
if args.no_save:
return
ckpt_dir = os.path.join(output_dir, 'ckpt')
if not os.path.exists(ckpt_dir):
os.makedirs(ckpt_dir)
save_name = os.path.join(ckpt_dir, 'model_{}_{}.pth'.format(args.epoch, args.step))
if isinstance(model, mynn.DataParallel):
model = model.module
# TODO: (maybe) Do not save redundant shared params
# model_state_dict = model.state_dict()
torch.save({
'epoch': args.epoch,
'step': args.step,
'iters_per_epoch': args.iters_per_epoch,
'model': model.state_dict(),
'optimizer': optimizer.state_dict()}, save_name)
logger.info('save model: %s', save_name) 示例2: initialize_model_from_cfg
# 需要导入模块: import nn [as 别名]
# 或者: from nn import DataParallel [as 别名]
def initialize_model_from_cfg(args, gpu_id=0):
"""Initialize a model from the global cfg. Loads test-time weights and
set to evaluation mode.
"""
model = model_builder.Generalized_RCNN()
model.eval()
if args.cuda:
model.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
logger.info("loading checkpoint %s", load_name)
checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage)
net_utils.load_ckpt(model, checkpoint['model'])
if args.load_detectron:
logger.info("loading detectron weights %s", args.load_detectron)
load_detectron_weight(model, args.load_detectron)
model = mynn.DataParallel(model, cpu_keywords=['im_info', 'roidb'], minibatch=True)
return model 示例3: save_ckpt
# 需要导入模块: import nn [as 别名]
# 或者: from nn import DataParallel [as 别名]
def save_ckpt(output_dir, args, step, train_size, model, optimizer):
"""Save checkpoint"""
if args.no_save:
return
ckpt_dir = os.path.join(output_dir, 'ckpt')
if not os.path.exists(ckpt_dir):
os.makedirs(ckpt_dir)
save_name = os.path.join(ckpt_dir, 'model_step{}.pth'.format(step))
if isinstance(model, mynn.DataParallel):
model = model.module
model_state_dict = model.state_dict()
torch.save({
'step': step,
'train_size': train_size,
'batch_size': args.batch_size,
'model': model.state_dict(),
'optimizer': optimizer.state_dict()}, save_name)
logger.info('save model: %s', save_name) 示例4: initialize_model_from_cfg
# 需要导入模块: import nn [as 别名]
# 或者: from nn import DataParallel [as 别名]
def initialize_model_from_cfg(args, gpu_id=0):
"""Initialize a model from the global cfg. Loads test-time weights and
set to evaluation mode.
"""
model = Generalized_RCNN()
model.eval()
if args.cuda:
model.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
logger.info("loading checkpoint %s", load_name)
checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage)
net_utils.load_ckpt(model, checkpoint['model'])
# model.load_state_dict(checkpoint['model'])
if args.load_detectron:
logger.info("loading detectron weights %s", args.load_detectron)
load_detectron_weight(model, args.load_detectron)
model = mynn.DataParallel(model, cpu_keywords=['im_info', 'roidb'], minibatch=True)
return model 示例5: initialize_model_from_cfg
# 需要导入模块: import nn [as 别名]
# 或者: from nn import DataParallel [as 别名]
def initialize_model_from_cfg(args, gpu_id=0):
"""Initialize a model from the global cfg. Loads test-time weights and
set to evaluation mode.
"""
model = model_builder_rel.Generalized_RCNN()
model.eval()
if args.cuda:
model.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
logger.info("loading checkpoint %s", load_name)
checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage)
net_utils.load_ckpt(model, checkpoint['model'])
if args.load_detectron:
logger.info("loading detectron weights %s", args.load_detectron)
load_detectron_weight(model, args.load_detectron)
model = mynn.DataParallel(model, cpu_keywords=['im_info', 'roidb'], minibatch=True)
return model