def train(args=VITConfig()):
    '\n    :param args:\n     - log_dir\n    '
    pretrained_model_config = CONFIGS[args.model_type]
    num_classes = (10 if (args.dataset == 'cifar10') else 100)
    model = VITransModel(pretrained_model_config, args.img_size, zero_head=True, num_classes=num_classes)
    model.load_from(np.load(args.pretrained_dir))
    model.to(device=args.device)
    num_params = count_parameters(model)
    if args.TB_log:
        os.makedirs(args.log_dir, exist_ok=True)
        writer = SummaryWriter(logdir=((args.log_dir + args.record_algo) + args.test_cycles))
    args.train_batch_size = (args.train_batch_size // args.gradient_accumulation_steps)
    (train_loader, test_loader) = get_loader(args)
    optimizer = torch.optim.SGD(model.parameters(), lr=args.learning_rate, momentum=0.9, weight_decay=args.weight_decay)
    t_total = args.num_steps
    if (args.decay_type == 'cosine'):
        scheduler = WarmupCosineSchedule(optimizer, warmup_steps=args.warmup_steps, t_total=t_total)
    else:
        scheduler = WarmupLinearSchedule(optimizer, warmup_steps=args.warmup_steps, t_total=t_total)
    model.zero_grad()
    set_seed(args.seed)
    losses = AverageMeter()
    global_step = 0
    best_acc = 0
    while True:
        model.train()
        epoch_iterator = tqdm(train_loader, desc='Training (X / X Steps) (loss=X.X)', bar_format='{l_bar}{r_bar}', dynamic_ncols=True)
        for (step, batch) in enumerate(epoch_iterator):
            batch = tuple((t.to(args.device) for t in batch))
            (x, y) = batch
            loss = model.forward(x, y)
            loss.backward()
            if (((step + 1) % args.gradient_accumulation_steps) == 0):
                losses.update((loss.item() * args.gradient_accumulation_steps))
                torch.nn.utils.clip_grad_norm(model.parameters(), 1.0)
                scheduler.step()
                optimizer.step()
                optimizer.zero_grad()
                global_step += 1
                epoch_iterator.set_description(('Training (%d / %d Steps) (loss=%2.5f)' % (global_step, t_total, losses.val)))
                writer.add_scalar('Train/loss', scalar_value=losses.val, global_step=global_step)
                writer.add_scalar('Train/lr', scalar_value=scheduler.get_lr()[0], global_step=global_step)
                if ((global_step % args.eval_every) == 0):
                    accuracy = valid(args, model, writer, test_loader, global_step)
                    if (best_acc < accuracy):
                        best_acc = accuracy
                        model_save((args.record_algo + args.test_cycles), model)
                    model.train()
                if ((global_step % t_total) == 0):
                    break
        losses.reset()
        if ((global_step % t_total) == 0):
            break
    writer.close()
    print(('===' * 30))
    print(('Best Accuracy: \t%f' % best_acc))
    print('End Training!')
    print(('===' * 30))
:param args:
 - log_dir
def train(args=VITConfig()):
    '\n    :param args:\n     - log_dir\n    '
    pretrained_model_config = CONFIGS[args.model_type]
    num_classes = (10 if (args.dataset == 'cifar10') else 100)
    model = VITransModel(pretrained_model_config, args.img_size, zero_head=True, num_classes=num_classes)
    model.load_from(np.load(args.pretrained_dir))
    model.to(device=args.device)
    num_params = count_parameters(model)
    if args.TB_log:
        os.makedirs(args.log_dir, exist_ok=True)
        writer = SummaryWriter(logdir=((args.log_dir + args.record_algo) + args.test_cycles))
    args.train_batch_size = (args.train_batch_size // args.gradient_accumulation_steps)
    (train_loader, test_loader) = get_loader(args)
    optimizer = torch.optim.SGD(model.parameters(), lr=args.learning_rate, momentum=0.9, weight_decay=args.weight_decay)
    t_total = args.num_steps
    if (args.decay_type == 'cosine'):
        scheduler = WarmupCosineSchedule(optimizer, warmup_steps=args.warmup_steps, t_total=t_total)
    else:
        scheduler = WarmupLinearSchedule(optimizer, warmup_steps=args.warmup_steps, t_total=t_total)
    model.zero_grad()
    set_seed(args.seed)
    losses = AverageMeter()
    global_step = 0
    best_acc = 0
    while True:
        model.train()
        epoch_iterator = tqdm(train_loader, desc='Training (X / X Steps) (loss=X.X)', bar_format='{l_bar}{r_bar}', dynamic_ncols=True)
        for (step, batch) in enumerate(epoch_iterator):
            batch = tuple((t.to(args.device) for t in batch))
            (x, y) = batch
            loss = model.forward(x, y)
            loss.backward()
            if (((step + 1) % args.gradient_accumulation_steps) == 0):
                losses.update((loss.item() * args.gradient_accumulation_steps))
                torch.nn.utils.clip_grad_norm(model.parameters(), 1.0)
                scheduler.step()
                optimizer.step()
                optimizer.zero_grad()
                global_step += 1
                epoch_iterator.set_description(('Training (%d / %d Steps) (loss=%2.5f)' % (global_step, t_total, losses.val)))
                writer.add_scalar('Train/loss', scalar_value=losses.val, global_step=global_step)
                writer.add_scalar('Train/lr', scalar_value=scheduler.get_lr()[0], global_step=global_step)
                if ((global_step % args.eval_every) == 0):
                    accuracy = valid(args, model, writer, test_loader, global_step)
                    if (best_acc < accuracy):
                        best_acc = accuracy
                        model_save((args.record_algo + args.test_cycles), model)
                    model.train()
                if ((global_step % t_total) == 0):
                    break
        losses.reset()
        if ((global_step % t_total) == 0):
            break
    writer.close()
    print(('===' * 30))
    print(('Best Accuracy: \t%f' % best_acc))
    print('End Training!')
    print(('===' * 30))