def get_model():
"""get model"""
model_lib = importlib.import_module(FLAGS.model)
model = model_lib.Model(FLAGS.num_classes)
if getattr(FLAGS, 'distributed', False):
gpu_id = init_dist()
if getattr(FLAGS, 'distributed_all_reduce', False):
model_wrapper = AllReduceDistributedDataParallel(model.cuda())
else:
model_wrapper = torch.nn.parallel.DistributedDataParallel(
model.cuda(), [gpu_id], gpu_id)
else:
model_wrapper = torch.nn.DataParallel(model).cuda()
return model, model_wrapper
def get_model():
"""get model"""
model_lib = importlib.import_module(FLAGS.model)
model = model_lib.Model(FLAGS.num_classes, input_size=FLAGS.image_size)
inputs = Variable(torch.randn((2, 3, 32, 32)))
output = model(inputs)
print(output)
print(output.size())
#input('pause')
if getattr(FLAGS, 'distributed', False):
gpu_id = init_dist()
if getattr(FLAGS, 'distributed_all_reduce', False):
# seems faster
model_wrapper = AllReduceDistributedDataParallel(model.cuda())
else:
model_wrapper = torch.nn.parallel.DistributedDataParallel(
model.cuda(), [gpu_id], gpu_id)
else:
model_wrapper = torch.nn.DataParallel(model).cuda()
return model, model_wrapper