src/paws_train.py [334:496]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - (loss, ploss, rloss, lr_stats), etime = gpu_timer(train_step) loss_meter.update(loss) ploss_meter.update(ploss) rloss_meter.update(rloss) time_meter.update(etime) if (itr % log_freq == 0) or np.isnan(loss) or np.isinf(loss): csv_logger.log(epoch + 1, itr, ploss_meter.avg, rloss_meter.avg, time_meter.avg) logger.info('[%d, %5d] loss: %.3f (%.3f %.3f) ' '(%d ms; %d ms)' % (epoch + 1, itr, loss_meter.avg, ploss_meter.avg, rloss_meter.avg, time_meter.avg, data_meter.avg)) if lr_stats is not None: logger.info('[%d, %5d] lr_stats: %.3f (%.2e, %.2e)' % (epoch + 1, itr, lr_stats.avg, lr_stats.min, lr_stats.max)) assert not np.isnan(loss), 'loss is nan' # -- logging/checkpointing logger.info('avg. loss %.3f' % loss_meter.avg) if rank == 0: save_dict = { 'encoder': encoder.state_dict(), 'opt': optimizer.state_dict(), 'epoch': epoch + 1, 'unlabel_prob': unlabeled_frac, 'loss': loss_meter.avg, 's_batch_size': s_batch_size, 'u_batch_size': u_batch_size, 'world_size': world_size, 'lr': lr, 'temperature': temperature, 'amp': scaler.state_dict() } torch.save(save_dict, latest_path) if best_loss is None or best_loss > loss_meter.avg: best_loss = loss_meter.avg logger.info('updating "best" checkpoint') torch.save(save_dict, best_path) if (epoch + 1) % checkpoint_freq == 0 \ or (epoch + 1) % 10 == 0 and epoch < checkpoint_freq: torch.save(save_dict, save_path.format(epoch=f'{epoch + 1}')) def load_checkpoint( r_path, encoder, opt, scaler, use_fp16=False ): checkpoint = torch.load(r_path, map_location='cpu') epoch = checkpoint['epoch'] # -- loading encoder encoder.load_state_dict(checkpoint['encoder']) logger.info(f'loaded encoder from epoch {epoch}') # -- loading optimizer opt.load_state_dict(checkpoint['opt']) if use_fp16: scaler.load_state_dict(checkpoint['amp']) logger.info(f'loaded optimizers from epoch {epoch}') logger.info(f'read-path: {r_path}') del checkpoint return encoder, opt, epoch def init_model( device, model_name='resnet50', use_pred=False, output_dim=128 ): if 'wide_resnet' in model_name: encoder = wide_resnet.__dict__[model_name](dropout_rate=0.0) hidden_dim = 128 else: encoder = resnet.__dict__[model_name]() hidden_dim = 2048 if 'w2' in model_name: hidden_dim *= 2 elif 'w4' in model_name: hidden_dim *= 4 # -- projection head encoder.fc = torch.nn.Sequential(OrderedDict([ ('fc1', torch.nn.Linear(hidden_dim, hidden_dim)), ('bn1', torch.nn.BatchNorm1d(hidden_dim)), ('relu1', torch.nn.ReLU(inplace=True)), ('fc2', torch.nn.Linear(hidden_dim, hidden_dim)), ('bn2', torch.nn.BatchNorm1d(hidden_dim)), ('relu2', torch.nn.ReLU(inplace=True)), ('fc3', torch.nn.Linear(hidden_dim, output_dim)) ])) # -- prediction head encoder.pred = None if use_pred: mx = 4 # 4x bottleneck prediction head pred_head = OrderedDict([]) pred_head['bn1'] = torch.nn.BatchNorm1d(output_dim) pred_head['fc1'] = torch.nn.Linear(output_dim, output_dim//mx) pred_head['bn2'] = torch.nn.BatchNorm1d(output_dim//mx) pred_head['relu'] = torch.nn.ReLU(inplace=True) pred_head['fc2'] = torch.nn.Linear(output_dim//mx, output_dim) encoder.pred = torch.nn.Sequential(pred_head) encoder.to(device) logger.info(encoder) return encoder def init_opt( encoder, iterations_per_epoch, start_lr, ref_lr, ref_mom, nesterov, warmup, num_epochs, weight_decay=1e-6, final_lr=0.0 ): param_groups = [ {'params': (p for n, p in encoder.named_parameters() if ('bias' not in n) and ('bn' not in n))}, {'params': (p for n, p in encoder.named_parameters() if ('bias' in n) or ('bn' in n)), 'LARS_exclude': True, 'weight_decay': 0} ] optimizer = SGD( param_groups, weight_decay=weight_decay, momentum=0.9, nesterov=nesterov, lr=ref_lr) scheduler = WarmupCosineSchedule( optimizer, warmup_steps=warmup*iterations_per_epoch, start_lr=start_lr, ref_lr=ref_lr, final_lr=final_lr, T_max=num_epochs*iterations_per_epoch) optimizer = LARS(optimizer, trust_coefficient=0.001) return encoder, optimizer, scheduler if __name__ == "__main__": main() - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - src/suncet_train.py [300:462]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - (loss, ploss, rloss, lr_stats), etime = gpu_timer(train_step) loss_meter.update(loss) ploss_meter.update(ploss) rloss_meter.update(rloss) time_meter.update(etime) if (itr % log_freq == 0) or np.isnan(loss) or np.isinf(loss): csv_logger.log(epoch + 1, itr, ploss_meter.avg, rloss_meter.avg, time_meter.avg) logger.info('[%d, %5d] loss: %.3f (%.3f %.3f) ' '(%d ms; %d ms)' % (epoch + 1, itr, loss_meter.avg, ploss_meter.avg, rloss_meter.avg, time_meter.avg, data_meter.avg)) if lr_stats is not None: logger.info('[%d, %5d] lr_stats: %.3f (%.2e, %.2e)' % (epoch + 1, itr, lr_stats.avg, lr_stats.min, lr_stats.max)) assert not np.isnan(loss), 'loss is nan' # -- logging/checkpointing logger.info('avg. loss %.3f' % loss_meter.avg) if rank == 0: save_dict = { 'encoder': encoder.state_dict(), 'opt': optimizer.state_dict(), 'epoch': epoch + 1, 'unlabel_prob': unlabeled_frac, 'loss': loss_meter.avg, 's_batch_size': s_batch_size, 'u_batch_size': u_batch_size, 'world_size': world_size, 'lr': lr, 'temperature': temperature, 'amp': scaler.state_dict() } torch.save(save_dict, latest_path) if best_loss is None or best_loss > loss_meter.avg: best_loss = loss_meter.avg logger.info('updating "best" checkpoint') torch.save(save_dict, best_path) if (epoch + 1) % checkpoint_freq == 0 \ or (epoch + 1) % 10 == 0 and epoch < checkpoint_freq: torch.save(save_dict, save_path.format(epoch=f'{epoch + 1}')) def load_checkpoint( r_path, encoder, opt, scaler, use_fp16=False ): checkpoint = torch.load(r_path, map_location='cpu') epoch = checkpoint['epoch'] # -- loading encoder encoder.load_state_dict(checkpoint['encoder']) logger.info(f'loaded encoder from epoch {epoch}') # -- loading optimizer opt.load_state_dict(checkpoint['opt']) if use_fp16: scaler.load_state_dict(checkpoint['amp']) logger.info(f'loaded optimizers from epoch {epoch}') logger.info(f'read-path: {r_path}') del checkpoint return encoder, opt, epoch def init_model( device, model_name='resnet50', use_pred=False, output_dim=128 ): if 'wide_resnet' in model_name: encoder = wide_resnet.__dict__[model_name](dropout_rate=0.0) hidden_dim = 128 else: encoder = resnet.__dict__[model_name]() hidden_dim = 2048 if 'w2' in model_name: hidden_dim *= 2 elif 'w4' in model_name: hidden_dim *= 4 # -- projection head encoder.fc = torch.nn.Sequential(OrderedDict([ ('fc1', torch.nn.Linear(hidden_dim, hidden_dim)), ('bn1', torch.nn.BatchNorm1d(hidden_dim)), ('relu1', torch.nn.ReLU(inplace=True)), ('fc2', torch.nn.Linear(hidden_dim, hidden_dim)), ('bn2', torch.nn.BatchNorm1d(hidden_dim)), ('relu2', torch.nn.ReLU(inplace=True)), ('fc3', torch.nn.Linear(hidden_dim, output_dim)) ])) # -- prediction head encoder.pred = None if use_pred: mx = 4 # 4x bottleneck prediction head pred_head = OrderedDict([]) pred_head['bn1'] = torch.nn.BatchNorm1d(output_dim) pred_head['fc1'] = torch.nn.Linear(output_dim, output_dim//mx) pred_head['bn2'] = torch.nn.BatchNorm1d(output_dim//mx) pred_head['relu'] = torch.nn.ReLU(inplace=True) pred_head['fc2'] = torch.nn.Linear(output_dim//mx, output_dim) encoder.pred = torch.nn.Sequential(pred_head) encoder.to(device) logger.info(encoder) return encoder def init_opt( encoder, iterations_per_epoch, start_lr, ref_lr, ref_mom, nesterov, warmup, num_epochs, weight_decay=1e-6, final_lr=0.0 ): param_groups = [ {'params': (p for n, p in encoder.named_parameters() if ('bias' not in n) and ('bn' not in n))}, {'params': (p for n, p in encoder.named_parameters() if ('bias' in n) or ('bn' in n)), 'LARS_exclude': True, 'weight_decay': 0} ] optimizer = SGD( param_groups, weight_decay=weight_decay, momentum=0.9, nesterov=nesterov, lr=ref_lr) scheduler = WarmupCosineSchedule( optimizer, warmup_steps=warmup*iterations_per_epoch, start_lr=start_lr, ref_lr=ref_lr, final_lr=final_lr, T_max=num_epochs*iterations_per_epoch) optimizer = LARS(optimizer, trust_coefficient=0.001) return encoder, optimizer, scheduler if __name__ == "__main__": main() - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -