def main(): parser = argparse.ArgumentParser(description=' FixMatch Training') parser.add_argument('--wresnet-k', default=2, type=int, help='width factor of wide resnet') parser.add_argument('--wresnet-n', default=28, type=int, help='depth of wide resnet') parser.add_argument('--dataset', type=str, default='CIFAR10', help='number of classes in dataset') # parser.add_argument('--n-classes', type=int, default=100, # help='number of classes in dataset') parser.add_argument('--n-labeled', type=int, default=40, help='number of labeled samples for training') parser.add_argument('--n-epoches', type=int, default=1024, help='number of training epoches') parser.add_argument('--batchsize', type=int, default=40, help='train batch size of labeled samples') parser.add_argument('--mu', type=int, default=7, help='factor of train batch size of unlabeled samples') parser.add_argument('--thr', type=float, default=0.95, help='pseudo label threshold') parser.add_argument('--n-imgs-per-epoch', type=int, default=64 * 1024, help='number of training images for each epoch') parser.add_argument('--lam-u', type=float, default=1., help='coefficient of unlabeled loss') parser.add_argument('--ema-alpha', type=float, default=0.999, help='decay rate for ema module') parser.add_argument('--lr', type=float, default=0.03, help='learning rate for training') parser.add_argument('--weight-decay', type=float, default=5e-4, help='weight decay') parser.add_argument('--momentum', type=float, default=0.9, help='momentum for optimizer') parser.add_argument('--seed', type=int, default=-1, help='seed for random behaviors, no seed if negtive') parser.add_argument('--temperature', type=float, default=0.5, help='temperature for loss function') args = parser.parse_args() # args.device = torch.device("cuda:2" if torch.cuda.is_available() else "cpu") logger, writer = setup_default_logging(args) logger.info(dict(args._get_kwargs())) # global settings # torch.multiprocessing.set_sharing_strategy('file_system') if args.seed > 0: torch.manual_seed(args.seed) random.seed(args.seed) np.random.seed(args.seed) # torch.backends.cudnn.deterministic = True n_iters_per_epoch = args.n_imgs_per_epoch // args.batchsize # 1024 n_iters_all = n_iters_per_epoch * args.n_epoches # 1024 * 1024 logger.info("***** Running training *****") logger.info(f" Task = {args.dataset}@{args.n_labeled}") logger.info(f" Num Epochs = {n_iters_per_epoch}") logger.info(f" Batch size per GPU = {args.batchsize}") # logger.info(f" Total train batch size = {args.batch_size * args.world_size}") logger.info(f" Total optimization steps = {n_iters_all}") model, criteria_x, criteria_u, criteria_z = set_model(args) logger.info("Total params: {:.2f}M".format( sum(p.numel() for p in model.parameters()) / 1e6)) dltrain_x, dltrain_u = get_train_loader(args.dataset, args.batchsize, args.mu, n_iters_per_epoch, L=args.n_labeled) dlval = get_val_loader(dataset=args.dataset, batch_size=64, num_workers=2) lb_guessor = LabelGuessor(thresh=args.thr) ema = EMA(model, args.ema_alpha) wd_params, non_wd_params = [], [] for name, param in model.named_parameters(): # if len(param.size()) == 1: if 'bn' in name: non_wd_params.append( param) # bn.weight, bn.bias and classifier.bias # print(name) else: wd_params.append(param) param_list = [{ 'params': wd_params }, { 'params': non_wd_params, 'weight_decay': 0 }] optim = torch.optim.SGD(param_list, lr=args.lr, weight_decay=args.weight_decay, momentum=args.momentum, nesterov=True) lr_schdlr = WarmupCosineLrScheduler(optim, max_iter=n_iters_all, warmup_iter=0) train_args = dict(model=model, criteria_x=criteria_x, criteria_u=criteria_u, criteria_z=criteria_z, optim=optim, lr_schdlr=lr_schdlr, ema=ema, dltrain_x=dltrain_x, dltrain_u=dltrain_u, lb_guessor=lb_guessor, lambda_u=args.lam_u, n_iters=n_iters_per_epoch, logger=logger) best_acc = -1 best_epoch = 0 logger.info('-----------start training--------------') for epoch in range(args.n_epoches): train_loss, loss_x, loss_u, loss_u_real, loss_simclr, mask_mean = train_one_epoch( epoch, **train_args) # torch.cuda.empty_cache() top1, top5, valid_loss = evaluate(ema, dlval, criteria_x) writer.add_scalars('train/1.loss', { 'train': train_loss, 'test': valid_loss }, epoch) writer.add_scalar('train/2.train_loss_x', loss_x, epoch) writer.add_scalar('train/3.train_loss_u', loss_u, epoch) writer.add_scalar('train/4.train_loss_u_real', loss_u_real, epoch) writer.add_scalar('train/4.train_loss_simclr', loss_simclr, epoch) writer.add_scalar('train/5.mask_mean', mask_mean, epoch) writer.add_scalars('test/1.test_acc', { 'top1': top1, 'top5': top5 }, epoch) # writer.add_scalar('test/2.test_loss', loss, epoch) # best_acc = top1 if best_acc < top1 else best_acc if best_acc < top1: best_acc = top1 best_epoch = epoch logger.info( "Epoch {}. Top1: {:.4f}. Top5: {:.4f}. best_acc: {:.4f} in epoch{}" .format(epoch, top1, top5, best_acc, best_epoch)) writer.close()
def main(): parser = argparse.ArgumentParser(description=' FixMatch Training') parser.add_argument('--wresnet-k', default=2, type=int, help='width factor of wide resnet') parser.add_argument('--wresnet-n', default=28, type=int, help='depth of wide resnet') parser.add_argument('--dataset', type=str, default='CIFAR10', help='number of classes in dataset') # parser.add_argument('--n-classes', type=int, default=100, # help='number of classes in dataset') parser.add_argument('--n-labeled', type=int, default=40, help='number of labeled samples for training') parser.add_argument('--n-epoches', type=int, default=1024, help='number of training epoches') parser.add_argument('--batchsize', type=int, default=40, help='train batch size of labeled samples') parser.add_argument('--mu', type=int, default=7, help='factor of train batch size of unlabeled samples') parser.add_argument('--thr', type=float, default=0.95, help='pseudo label threshold') parser.add_argument('--n-imgs-per-epoch', type=int, default=64 * 1024, help='number of training images for each epoch') parser.add_argument('--lam-u', type=float, default=1., help='coefficient of unlabeled loss') parser.add_argument('--lam-s', type=float, default=0.2, help='coefficient of unlabeled loss SimCLR') parser.add_argument('--ema-alpha', type=float, default=0.999, help='decay rate for ema module') parser.add_argument('--lr', type=float, default=0.03, help='learning rate for training') parser.add_argument('--weight-decay', type=float, default=5e-4, help='weight decay') parser.add_argument('--momentum', type=float, default=0.9, help='momentum for optimizer') parser.add_argument('--seed', type=int, default=-1, help='seed for random behaviors, no seed if negtive') parser.add_argument('--temperature', type=float, default=0.5, help='temperature for loss function') args = parser.parse_args() # args.device = torch.device("cuda:2" if torch.cuda.is_available() else "cpu") logger, writer = setup_default_logging(args) logger.info(dict(args._get_kwargs())) # global settings # torch.multiprocessing.set_sharing_strategy('file_system') if args.seed > 0: torch.manual_seed(args.seed) random.seed(args.seed) np.random.seed(args.seed) # torch.backends.cudnn.deterministic = True n_iters_per_epoch = args.n_imgs_per_epoch // args.batchsize # 1024 n_iters_all = n_iters_per_epoch * args.n_epoches # 1024 * 1024 logger.info("***** Running training *****") logger.info(f" Task = {args.dataset}@{args.n_labeled}") logger.info(f" Num Epochs = {n_iters_per_epoch}") logger.info(f" Batch size per GPU = {args.batchsize}") # logger.info(f" Total train batch size = {args.batch_size * args.world_size}") logger.info(f" Total optimization steps = {n_iters_all}") model, criteria_x, criteria_u, criteria_z = set_model(args) logger.info("Total params: {:.2f}M".format( sum(p.numel() for p in model.parameters()) / 1e6)) dltrain_x, dltrain_u, dltrain_f = get_train_loader_mix(args.dataset, args.batchsize, args.mu, n_iters_per_epoch, L=args.n_labeled) dlval = get_val_loader(dataset=args.dataset, batch_size=64, num_workers=2) lb_guessor = LabelGuessor(thresh=args.thr) ema = EMA(model, args.ema_alpha) wd_params, non_wd_params = [], [] for name, param in model.named_parameters(): # if len(param.size()) == 1: if 'bn' in name: non_wd_params.append( param) # bn.weight, bn.bias and classifier.bias # print(name) else: wd_params.append(param) param_list = [{ 'params': wd_params }, { 'params': non_wd_params, 'weight_decay': 0 }] optim_fix = torch.optim.SGD(param_list, lr=args.lr, weight_decay=args.weight_decay, momentum=args.momentum, nesterov=True) lr_schdlr_fix = WarmupCosineLrScheduler(optim_fix, max_iter=n_iters_all, warmup_iter=0) train_args = dict(model=model, criteria_x=criteria_x, criteria_u=criteria_u, criteria_z=criteria_z, optim=optim_fix, lr_schdlr=lr_schdlr_fix, ema=ema, dltrain_x=dltrain_x, dltrain_u=dltrain_u, dltrain_f=dltrain_f, lb_guessor=lb_guessor, lambda_u=args.lam_u, lambda_s=args.lam_s, n_iters=n_iters_per_epoch, logger=logger, bt=args.batchsize, mu=args.mu) # # TRAINING PARAMETERS FOR SIMCLR # param_list = [ # {'params': wd_params}, {'params': non_wd_params, 'weight_decay': 0}] # # optim_simclr = torch.optim.SGD(param_list, lr=0.5, weight_decay=args.weight_decay, # momentum=args.momentum, nesterov=False) # # lr_schdlr_simclr = WarmupCosineLrScheduler( # optim_simclr, max_iter=n_iters_all, warmup_iter=0 # ) # # train_args_simclr = dict( # model=model, # criteria_z=criteria_z, # optim=optim_simclr, # lr_schdlr=lr_schdlr_simclr, # ema=ema, # dltrain_f=dltrain_f, # lambda_s=args.lam_s, # n_iters=n_iters_per_epoch, # logger=logger, # bt=args.batchsize, # mu=args.mu # ) # # TRAINING PARAMETERS FOR IIC # param_list = [ # {'params': wd_params}, {'params': non_wd_params, 'weight_decay': 0}] # # optim_iic = torch.optim.Adam(param_list, lr=1e-4, weight_decay=args.weight_decay) # # lr_schdlr_iic = WarmupCosineLrScheduler( # optim_iic, max_iter=n_iters_all, warmup_iter=0 # ) # # train_args_iic = dict( # model=model, # optim=optim_iic, # lr_schdlr=lr_schdlr_iic, # ema=ema, # dltrain_f=dltrain_f, # n_iters=n_iters_per_epoch, # logger=logger, # bt=args.batchsize, # mu=args.mu # ) # best_acc = -1 best_epoch = 0 logger.info('-----------start training--------------') for epoch in range(args.n_epoches): # guardar accuracy de modelo preentrenado hasta espacio h (SALIDA DE BACKBONE) top1, top5, valid_loss = evaluate_linear_Clf(ema, dltrain_x, dlval, criteria_x) writer.add_scalars('test/1.test_linear_acc', { 'top1': top1, 'top5': top5 }, epoch) logger.info("Epoch {}. on h space Top1: {:.4f}. Top5: {:.4f}.".format( epoch, top1, top5)) if epoch < -500: # # FASE DE ENTRENAMIENTO NO SUPERVISADO # entrenar feature representation simclr # train_loss, loss_simclr, model_ = train_one_epoch_simclr(epoch, **train_args_simclr) # writer.add_scalar('train/4.train_loss_simclr', loss_simclr, epoch) # entrenar iic # train_loss, loss_iic, model_ = train_one_epoch_iic(epoch, **train_args_iic) # writer.add_scalar('train/4.train_loss_iic', loss_iic, epoch) # evaluate_Clf(model_, dltrain_f, dlval, criteria_x) top1, top5, valid_loss = evaluate_linear_Clf( ema, dltrain_x, dlval, criteria_x) # # GUARDAR MODELO ENTRENADO DE FORMA NO SUPERVISADA # if epoch == 497: # # save model # name = 'simclr_trained_good_h2.pt' # torch.save(model_.state_dict(), name) # logger.info('model saved') else: # ENTRENAMIENTO SEMI-SUPERVISADO train_loss, loss_x, loss_u, loss_u_real, mask_mean, loss_simclr = train_one_epoch( epoch, **train_args) top1, top5, valid_loss = evaluate(ema, dlval, criteria_x) writer.add_scalar('train/4.train_loss_simclr', loss_simclr, epoch) writer.add_scalar('train/2.train_loss_x', loss_x, epoch) writer.add_scalar('train/3.train_loss_u', loss_u, epoch) writer.add_scalar('train/4.train_loss_u_real', loss_u_real, epoch) writer.add_scalar('train/5.mask_mean', mask_mean, epoch) writer.add_scalars('train/1.loss', { 'train': train_loss, 'test': valid_loss }, epoch) writer.add_scalars('test/1.test_acc', { 'top1': top1, 'top5': top5 }, epoch) # writer.add_scalar('test/2.test_loss', loss, epoch) # best_acc = top1 if best_acc < top1 else best_acc if best_acc < top1: best_acc = top1 best_epoch = epoch logger.info( "Epoch {}. Top1: {:.4f}. Top5: {:.4f}. best_acc: {:.4f} in epoch{}" .format(epoch, top1, top5, best_acc, best_epoch)) writer.close()