def init_model(args, field, logger, world_size, device):
logger.info(f'Initializing {args.model}')
Model = getattr(models, args.model)
model = Model(field, args)
params = get_trainable_params(model)
num_param = count_params(params)
logger.info(f'{args.model} has {num_param:,} trainable parameters')
model.to(device)
if world_size > 1:
logger.info(f'Wrapping model for distributed')
model = DistributedDataParallel(model)
model.params = params
return model
def init_model(args, field, logger, world_size):
logger.info(f'Initializing {args.model}')
Model = getattr(models, args.model)
model = Model(field, args)
# 模型初始化
params = get_trainable_params(model)
num_param = count_params(params)
# 计算模型参数个数
logger.info(f'{args.model} has {num_param:,} parameters')
if args.gpus[0] > -1:
model.cuda()
# 是否使用gpu设置的地方，如果设置为-1或者更负，就不使用gpu，只用cpu
if world_size > 1:
logger.info(f'Wrapping model for distributed')
model = DistributedDataParallel(model)
model.params = params
return model