"""

3种运行方式：

1. 单CPU运行模式；

2. 单GPU运行模式；

3. 分布式运行模式：多机多卡 或 单机多卡。

分布式优势：1.支持同步BN； 2.DDP每个训练有独立进程管理，训练速度更快，显存均衡；

"""

args = parser.parse_args()

# 根据训练机器和超参，选择运行方式

num_gpus_available = torch.cuda.device_count()

if num_gpus_available == 0:

args.gpus = 0

elif args.gpus > num_gpus_available:

raise ValueError(f'--gpus(-g {args.gpus}) can not greater than available device({num_gpus_available})')

# 根据每个节点的GPU数量调整world size

args.world_size = args.gpus * args.nodes

if not args.cuda or args.world_size == 0:

# 1. cpu运行模式

args.cuda = False

args.gpus = 0

args.distributed = False

elif args.world_size == 1:

# 2. 单GPU运行模式

args.distributed = False

elif args.world_size > 1:

# 3. 分布式运行模式

args.distributed = True

else:

raise ValueError(f'Check config parameters --nodes/-n={args.nodes} and --gpus/-g={args.gpus}!')

if args.distributed and args.gpus > 1:

# use torch.multiprocessing.spawn to launch distributed processes

mp.spawn(main_worker, nprocs=args.gpus, args=(args,))

else:

# Simply call main_worker function

""" 模型训练、测试、转JIT、蒸馏文件制作

:param gpu: 运行的gpu id

:param args: 运行超参

"""

args.gpu = gpu

utils.generate_logger(f"{datetime.datetime.now().strftime('%Y%m%d%H%M%S')}-{gpu}.log")

logging.info(f'args: {args}')

# 可复现性

if args.seed is not None:

random.seed(args.seed)

torch.manual_seed(args.seed)

cudnn.deterministic = True

logging.warning('You have chosen to seed training. '

'This will turn on the CUDNN deterministic setting, '

'which can slow down your training considerably! '

'You may see unexpected behavior when restarting '

'from checkpoints.')

if args.cuda:

logging.info(f"Use GPU: {args.gpu} ~")

if args.distributed:

args.rank = args.rank * args.gpus + gpu

dist.init_process_group(backend='nccl', init_method=args.init_method,

world_size=args.world_size, rank=args.rank)

else:

logging.info(f"Use CPU ~")

# 创建/加载模型，使用预训练模型时，需要自己先下载好放到 pretrained 文件夹下，以网络名词命名

logging.info(f"=> creating model '{args.arch}'")

model = my_models.get_model(args.arch, args.pretrained, num_classes=args.num_classes)

# 重加载之前训练好的模型

if args.resume:

if os.path.isfile(args.resume):

logging.info(f"=> loading checkpoint '{args.resume}'")

checkpoint = torch.load(args.resume, map_location=torch.device('cpu'))

acc = model.load_state_dict(checkpoint['state_dict'])

logging.info(f'missing keys of models: {acc.missing_keys}')

del checkpoint

else:

raise Exception(f"No checkpoint found at '{args.resume}' to be resumed")

# 模型信息

image_height, image_width = args.image_size

logging.info(f'Model {args.arch} input size: ({image_height}, {image_width})')

utils.summary(size=(image_height, image_width), channel=3, model=model)

# 模型转换：转为 torch.jit.script

if args.jit:

if not args.resume:

raise Exception('Option --resume must specified!')

applications.convert_to_jit(model, args=args)

return

if args.criterion == 'rank':

criterion = criterions.RankingLoss(args=args) # 对比排序损失

elif args.criterion == 'emd':

criterion = criterions.EMDLoss() # 推土机距离损失

elif args.criterion == 'regress':

criterion = criterions.RegressionLoss() # MSE回归损失

else:

raise NotImplementedError(f'Not loss function {args.criterion}，only (rank, emd, regress)!')

if args.cuda:

if args.distributed and args.sync_bn:

model = apex.parallel.convert_syncbn_model(model)

torch.cuda.set_device(args.gpu)

model.cuda(args.gpu)

criterion = criterion.cuda(args.gpu)

# 优化器：Adam > SGD > SWA(SGD > Adam)

optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)

# 可尝试优化器

# optimizer = torch.optim.SGD(model.parameters(),

# args.lr, momentum=args.momentum,

# weight_decay=args.weight_decay)

# from optim.torchtools.optim import RangerLars, Ralamb, Novograd, LookaheadAdam, Ranger, RAdam, AdamW

# optimizer = RangerLars(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)

# optimizer = Ralamb(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)

# optimizer = Novograd(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)

# optimizer = LookaheadAdam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)

# optimizer = Ranger(model_params, lr=args.lr, weight_decay=args.weight_decay)

# optimizer = RAdam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)

# optimizer = AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)

# 随机均值平均优化器

# from optim.swa import SWA

# optimizer = SWA(optimizer, swa_start=10, swa_freq=5, swa_lr=0.05)

# 混合精度训练

if args.cuda:

model, optimizer = amp.initialize(model, optimizer, opt_level="O1")

model = DDP(model)

else:

model = torch.nn.DataParallel(model)

if args.train:

train_loader = dataloader.load(args, 'train')

val_loader = dataloader.load(args, 'val')

scheduler = LambdaLR(optimizer,

lambda epoch: adjust_learning_rate(epoch, args=args))

applications.train(train_loader, val_loader, model, criterion, optimizer, scheduler, args)

args.evaluate = True

if args.evaluate:

torch.set_flush_denormal(True)

test_loader = dataloader.load(args, name='test')

acc, loss, test_results = applications.test(test_loader, model, criterion, args)

logging.info(f'Evaluation: * Acc@1 {acc:.3f} and loss {loss:.3f}.')

logging.info(f'Evaluation results:')

for result in test_results:

logging.info(' '.join([str(r) for r in result]))

""" 根据warmup设置、迭代代数、设置的学习率，调整每一代的学习率

:param epoch: 当前epoch数

:param args: 使用warmup代数、学习率

"""

# lr_rates = [0.1, 1., 10., 100., 1e-10]

# epochs = [2, 4, 6, 8, 10]

epoch_step = (args.epochs - args.warmup) / 4.0

lr_rates = np.array([0.1, 1., 0.1, 0.01, 0.001])

epochs = np.array([args.warmup,

args.warmup + int(1.5 * epoch_step),

args.warmup + int(2.5 * epoch_step),

args.warmup + int(3.5 * epoch_step),

args.epochs])

for i, e in enumerate(epochs):

if e > epoch:

return lr_rates[i]

elif e == epoch:

next_rate = lr_rates[i]

if len(lr_rates) > i + 1:

next_rate = lr_rates[i + 1]

logging.info(f'===== lr decay rate: {lr_rates[i]} -> {next_rate} =====')