# Do we have a config file to parse?

args_config, remaining = config_parser.parse_known_args()

if args_config.config:

with open(args_config.config, 'r') as f:

cfg = yaml.safe_load(f)

parser.set_defaults(**cfg)

# The main arg parser parses the rest of the args, the usual

# defaults will have been overridden if config file specified.

args = parser.parse_args(remaining)

if args.qmodules is None:

args.qmodules = []

# Cache the args as a text string to save them in the output dir later

args_text = yaml.safe_dump(args.__dict__, default_flow_style=False)

qat_model = copy.deepcopy(model)

qat_model.train()

qconfigs = {}

ACT_LAYER_MAPPINGS = {

'relu': nn.ReLU,

'gelu': nn.GELU,

'prelu' : nn.PReLU,

'rprelu' : 'rprelu',

'None': 'None'

}

for m in args.qmodules:

wcfg = {

"mode": args.wq_mode if args.wq_enable else "Identity",

"bit": args.wq_bitw if args.wq_bitw < 32 and args.aq_enable else "identity",

"all_positive": False,

"symmetric": not args.wq_asym,

"per_channel": args.wq_per_channel,

"normalize_first": False,

"learnable": args.wq_clip_learnable

}

acfg = {

"enable": args.aq_enable if args.aq_enable else "Identity",

"mode": args.aq_mode if args.aq_bitw < 32 and args.aq_enable else "identity",

"bit": args.aq_bitw,

"per_channel": args.aq_per_channel,

"normalize_first": False,

"learnable": args.aq_clip_learnable

}

qconfigs[m] = {"weight": wcfg, "act": acfg, "q_attn_dropout": args.apply_q_attn_dropout, "act_layer": ACT_LAYER_MAPPINGS[args.act_layer] }

if args.model_type == 'deit':

qat_model = replace_module_by_qmodule_deit(model, qconfigs, pretrained_initialized = args.pretrained_initialized, \

qk_reparam=args.qk_reparam, qk_reparam_type = args.qk_reparam_type)

elif args.model_type == 'swin':

qat_model = replace_module_by_qmodule_swin(model, qconfigs, pretrained_initialized = args.pretrained_initialized, \

qk_reparam=args.qk_reparam, qk_reparam_type = args.qk_reparam_type)

if args.teacher_type == "deit":

teacher = create_model(args.teacher,num_classes=args.num_classes,

drop_rate=args.drop, pretrained=args.teacher_pretrained, qqkkvv = args.kd_hard_and_soft == 2 or args.kd_hard_and_soft == 3

)

elif args.teacher_type == "swin":

teacher = create_model(args.teacher, num_classes=args.num_classes, drop_path=args.drop_path, pretrained=args.teacher_pretrained, qqkkvv = args.kd_hard_and_soft == 2 or args.kd_hard_and_soft == 3)

if args.quant_teacher:

args.wq_bitw = 4

args.aq_bitw = 4

teacher = get_qat_model(teacher, args)

if args.teacher_pretrained and args.teacher_checkpoint != "":

load_checkpoint(teacher, args.teacher_checkpoint, strict=True)

setup_default_logging()

args, args_text = args

args.use_kd = args.use_kd or args.use_token_kd

args.prefetcher = not args.no_prefetcher

args.local_rank = local_rank

if args.log_wandb:

if has_wandb:

if args.local_rank == 0:

print("init wandb")

wandb.init(project=args.experiment, config=args)

else:

_logger.warning("You've requested to log metrics to wandb but package not found. "

"Metrics not being logged to wandb, try `pip install wandb`")

args.distributed = False

if 'WORLD_SIZE' in os.environ:

args.distributed = int(os.environ['WORLD_SIZE']) > 1

print("args.local_rank: "+str(args.local_rank))

print(int(os.environ['WORLD_SIZE']))

args.device = 'cuda:'+str(args.gpu_id)

args.world_size = 1

args.rank = 0 # global rank

if args.distributed:

torch.distributed.init_process_group(backend='nccl', init_method='tcp://localhost:'+args.tcp_port,rank=args.local_rank, world_size=int(os.environ['WORLD_SIZE']))

args.device = 'cuda:%d' % args.local_rank

torch.cuda.set_device(args.local_rank)

args.world_size = torch.distributed.get_world_size()

args.rank = torch.distributed.get_rank()

_logger.info('Training in distributed mode with multiple processes, 1 GPU per process. Process %d, total %d.'

% (args.rank, args.world_size))

else:

_logger.info('Training with a single process on 1 GPUs.')

assert args.rank >= 0

# resolve AMP arguments based on PyTorch / Apex availability

use_amp = None

if args.amp:

# `--amp` chooses native amp before apex (APEX ver not actively maintained)

if has_native_amp:

args.native_amp = True

elif has_apex:

args.apex_amp = True

if args.apex_amp and has_apex:

use_amp = 'apex'

elif args.native_amp and has_native_amp:

use_amp = 'native'

elif args.apex_amp or args.native_amp:

_logger.warning("Neither APEX or native Torch AMP is available, using float32. "

"Install NVIDA apex or upgrade to PyTorch 1.6")

random_seed(args.seed, args.rank)

if args.model_type == "deit":

model = create_model(args.model,num_classes=args.num_classes,

drop_rate=args.drop, pretrained=args.pretrained, qqkkvv = args.kd_hard_and_soft == 2 or args.kd_hard_and_soft == 3

)

elif args.model_type == "swin":

model = create_model(args.model, drop_path=args.drop_path, num_classes=args.num_classes, pretrained=args.pretrained,

qqkkvv = args.kd_hard_and_soft == 2 or args.kd_hard_and_soft == 3 )

if args.quantized:

model = get_qat_model(model, args)

if args.initial_checkpoint != "":

load_checkpoint(model, args.initial_checkpoint, strict=False)

if args.num_classes is None:

assert hasattr(model, 'num_classes'), 'Model must have `num_classes` attr if not set on cmd line/config.'

args.num_classes = model.num_classes # FIXME handle model default vs config num_classes more elegantly

if args.replace_ln_by_bn:

model = replace_ln_by_bn1d(model)

# distillation

teacher = None

if args.use_kd:

if args.local_rank == 0:

print("create teacher model")

teacher = create_teacher_model(args)

teacher.cuda()

if args.local_rank == 0:

_logger.info(

f'Model {safe_model_name(args.model)} created, param count:{sum([m.numel() for m in model.parameters()])}')

data_config = resolve_data_config(vars(args), model=model, verbose=args.local_rank == 0)

# setup augmentation batch splits for contrastive loss or split bn

num_aug_splits = 0

if args.aug_splits > 0:

assert args.aug_splits > 1, 'A split of 1 makes no sense'

num_aug_splits = args.aug_splits

# enable split bn (separate bn stats per batch-portion)

if args.split_bn:

assert num_aug_splits > 1 or args.resplit

model = convert_splitbn_model(model, max(num_aug_splits, 2))

# move model to GPU, enable channels last layout if set

model.cuda()

if args.channels_last:

model = model.to(memory_format=torch.channels_last)

# setup synchronized BatchNorm for distributed training

if args.distributed and args.sync_bn:

assert not args.split_bn

if has_apex and use_amp != 'native':

# Apex SyncBN preferred unless native amp is activated

model = convert_syncbn_model(model)

else:

model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)

if args.local_rank == 0:

_logger.info(

'Converted model to use Synchronized BatchNorm. WARNING: You may have issues if using '

'zero initialized BN layers (enabled by default for ResNets) while sync-bn enabled.')

if args.torchscript:

assert not use_amp == 'apex', 'Cannot use APEX AMP with torchscripted model'

assert not args.sync_bn, 'Cannot use SyncBatchNorm with torchscripted model'

model = torch.jit.script(model)

## create datasets

dataset_train = create_dataset(

args.dataset,

root=args.data_dir, split=args.train_split, is_training=True,

batch_size=args.batch_size, repeats=args.epoch_repeats)

# setup mixup / cutmix

collate_fn = None

mixup_fn = None

mixup_active = args.mixup > 0 or args.cutmix > 0. or args.cutmix_minmax is not None

if mixup_active:

mixup_args = dict(

mixup_alpha=args.mixup, cutmix_alpha=args.cutmix, cutmix_minmax=args.cutmix_minmax,

prob=args.mixup_prob, switch_prob=args.mixup_switch_prob, mode=args.mixup_mode,

label_smoothing=args.smoothing, num_classes=args.num_classes)

if args.prefetcher:

assert not num_aug_splits # collate conflict (need to support deinterleaving in collate mixup)

collate_fn = FastCollateMixup(**mixup_args)

else:

mixup_fn = Mixup(**mixup_args)

# wrap dataset in AugMix helper

if num_aug_splits > 1:

dataset_train = AugMixDataset(dataset_train, num_splits=num_aug_splits)

# create data loaders w/ augmentation pipeiine

train_interpolation = args.train_interpolation

if args.no_aug or not train_interpolation:

train_interpolation = data_config['interpolation']

loader_train = create_loader(

dataset_train,

input_size=data_config['input_size'],

batch_size=args.batch_size,

is_training=True,

use_prefetcher=args.prefetcher,

no_aug=args.no_aug,

re_prob=args.reprob,

re_mode=args.remode,

re_count=args.recount,

re_split=args.resplit,

scale=args.scale,

ratio=args.ratio,

hflip=args.hflip,

vflip=args.vflip,

color_jitter=args.color_jitter,

auto_augment=args.aa,

num_aug_splits=num_aug_splits,

num_aug_repeats = args.num_aug_repeats,

interpolation=train_interpolation,

mean=data_config['mean'],

std=data_config['std'],

num_workers=args.workers,

distributed=args.distributed,

collate_fn=collate_fn,

pin_memory=args.pin_mem,

use_multi_epochs_loader=args.use_multi_epochs_loader

)

if args.local_rank == 0:

_logger.info(f"Trainer transform: {loader_train.dataset.transform}")

dataset_eval = create_dataset(

args.dataset, root=args.data_dir, split=args.val_split, is_training=False, batch_size=args.batch_size)

loader_eval = create_loader(

dataset_eval,

input_size=data_config['input_size'],

batch_size=args.validation_batch_size_multiplier * args.batch_size,

is_training=False,

use_prefetcher=args.prefetcher,

interpolation=data_config['interpolation'],

mean=data_config['mean'],

std=data_config['std'],

num_workers=args.workers,

distributed=args.distributed,

crop_pct=data_config['crop_pct'],

pin_memory=args.pin_mem,

)

print(dataset_train.__len__())

print(dataset_eval.__len__())

## have to create and init all the parameters before send to optimizer (important!)

setup_alpha(model=model,loader=loader_train,args=args)

if args.local_rank == 0:

_logger.info(f"Model {model}")

# when creating optimizer, only weights and emb have weight-decay

optimizer = create_optimizer_v2(model, **optimizer_kwargs(cfg=args))

# setup automatic mixed-precision (AMP) loss scaling and op casting

amp_autocast = suppress # do nothing

loss_scaler = None

if use_amp == 'apex':

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

loss_scaler = ApexScaler()

if args.local_rank == 0:

_logger.info('Using NVIDIA APEX AMP. Training in mixed precision.')

elif use_amp == 'native':

amp_autocast = torch.cuda.amp.autocast

loss_scaler = NativeScaler()

if args.local_rank == 0:

_logger.info('Using native Torch AMP. Training in mixed precision.')

else:

if args.local_rank == 0:

_logger.info('AMP not enabled. Training in float32.')

# optionally resume from a checkpoint

resume_epoch = None

output_dir_resume = None

if args.experiment:

try:

output_dir_resume = os.path.join(args.output, args.experiment)

except:

print('please indicate args.resume if you would like to resume the checkpoint')

if args.resume:

resume_epoch = resume_checkpoint(

model, args.resume,

optimizer=None if args.no_resume_opt else optimizer,

loss_scaler=None if args.no_resume_opt else loss_scaler,

log_info=args.local_rank == 0)

elif args.experiment:

if os.path.exists(os.path.join(output_dir_resume,'last.pth.tar')):

print("resume from last.pth.tar")

resume_path = os.path.join(output_dir_resume,'last.pth.tar')

resume_epoch = resume_checkpoint(

model, resume_path,

optimizer=None if args.no_resume_opt else optimizer,

loss_scaler=None if args.no_resume_opt else loss_scaler,

log_info=args.local_rank == 0)

# setup exponential moving average of model weights, SWA could be used here too

model_ema = None

if args.model_ema:

# Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper

model_ema = ModelEmaV2(

model, decay=args.model_ema_decay, device='cpu' if args.model_ema_force_cpu else None)

if args.resume:

load_checkpoint(model_ema.module, args.resume, use_ema=True)

# setup distributed training

if args.distributed:

if has_apex and use_amp != 'native':

# Apex DDP preferred unless native amp is activated

if args.local_rank == 0:

_logger.info("Using NVIDIA APEX DistributedDataParallel.")

model = ApexDDP(model, delay_allreduce=True)

else:

if args.local_rank == 0:

_logger.info("Using native Torch DistributedDataParallel.")

model = NativeDDP(model, device_ids=[args.local_rank]) # can use device str in Torch >= 1.1

# NOTE: EMA model does not need to be wrapped by DDP

# setup learning rate schedule and starting epoch

lr_scheduler, num_epochs = create_scheduler(args, optimizer)

start_epoch = 0

if args.start_epoch is not None:

# a specified start_epoch will always override the resume epoch

start_epoch = args.start_epoch

elif resume_epoch is not None:

start_epoch = resume_epoch

if lr_scheduler is not None and start_epoch > 0:

lr_scheduler.step(start_epoch)

if args.local_rank == 0:

_logger.info('Scheduled epochs: {}'.format(num_epochs))

print(args.dataset)

# setup loss function

if args.jsd:

assert num_aug_splits > 1 # JSD only valid with aug splits set

train_loss_fn = JsdCrossEntropy(num_splits=num_aug_splits, smoothing=args.smoothing).cuda()

elif args.use_token_kd:

train_loss_fn = KLTokenMSELoss(alpha=args.kd_alpha, kd_type=args.kd_type)

elif args.use_kd:

if args.kd_hard_and_soft == 0:

train_loss_fn = KLLossSoft()

elif args.kd_hard_and_soft == 1: ## this loss function will automatically apply soft cross entropy if the label probability is given

train_loss_fn = KDLossSoftandHard().cuda() # work with mixup as well

elif args.kd_hard_and_soft == 2:

train_loss_fn = KDLossSoftandHard_qk().cuda()

elif args.kd_hard_and_soft == 3:

train_loss_fn = KDLossSoftandHard_qkv().cuda()

elif mixup_active:

# smoothing is handled with mixup target transform

train_loss_fn = SoftTargetCrossEntropy().cuda()

elif args.smoothing:

train_loss_fn = LabelSmoothingCrossEntropy(smoothing=args.smoothing).cuda()

else:

train_loss_fn = nn.CrossEntropyLoss().cuda()

if args.local_rank == 0:

_logger.info(f"Train loss {train_loss_fn}")

validate_loss_fn = nn.CrossEntropyLoss().cuda()

if args.resume:

validate(model, loader_eval, validate_loss_fn, args, amp_autocast=amp_autocast,model_type=args.model_type)

elif args.experiment:

if os.path.exists(os.path.join(output_dir_resume,'last.pth.tar')):

validate(model, loader_eval, validate_loss_fn, args, amp_autocast=amp_autocast,model_type=args.model_type)

if args.log_wandb and has_wandb:

model.eval()

print("setup alpha")

with torch.no_grad():

for batch_idx, (input, target) in enumerate(loader):

if not args.prefetcher:

input = input.cuda()

target = target.cuda()

if args.channels_last:

input = input.contiguous(memory_format=torch.channels_last)

with amp_autocast():

output = model(input)

batch_time_m = AverageMeter()

losses_m = AverageMeter()

top1_m = AverageMeter()

top5_m = AverageMeter()

model.eval()

print("model eval")

end = time.time()

last_idx = len(loader) - 1

with torch.no_grad():

for batch_idx, (input, target) in enumerate(loader):

last_batch = batch_idx == last_idx

if not args.prefetcher:

input = input.cuda()

target = target.cuda()

if args.channels_last:

input = input.contiguous(memory_format=torch.channels_last)

with amp_autocast():

if model_type == 'rednet':

output = model.forward_test(input, softmax = False, post_process=False)

else:

output = model(input)

if isinstance(output, (tuple, list)):

output = output[0]

# augmentation reduction

reduce_factor = args.tta

if reduce_factor > 1:

output = output.unfold(0, reduce_factor, reduce_factor).mean(dim=2)

target = target[0:target.size(0):reduce_factor]

loss = loss_fn(output, target)

acc1, acc5 = accuracy(output, target, topk=(1, 5))

if args.distributed:

reduced_loss = reduce_tensor(loss.data, args.world_size)

acc1 = reduce_tensor(acc1, args.world_size)

acc5 = reduce_tensor(acc5, args.world_size)

else:

reduced_loss = loss.data

torch.cuda.synchronize()

losses_m.update(reduced_loss.item(), input.size(0))

top1_m.update(acc1.item(), output.size(0))

top5_m.update(acc5.item(), output.size(0))

batch_time_m.update(time.time() - end)

end = time.time()

if args.local_rank == 0 and (last_batch or batch_idx % args.log_interval == 0):

log_name = 'Test' + log_suffix

_logger.info(

'{0}: [{1:>4d}/{2}] '

'Time: {batch_time.val:.3f} ({batch_time.avg:.3f}) '

'Loss: {loss.val:>7.4f} ({loss.avg:>6.4f}) '

'Acc@1: {top1.val:>7.4f} ({top1.avg:>7.4f}) '

'Acc@5: {top5.val:>7.4f} ({top5.avg:>7.4f})'.format(

log_name, batch_idx, last_idx, batch_time=batch_time_m,

loss=losses_m, top1=top1_m, top5=top5_m))

if args.local_rank == 0 and log_wandb:

wandb.log({'val_loss': losses_m.avg, 'epoch': epoch})

wandb.log({'val_top1': top1_m.avg, 'epoch': epoch})

wandb.log({'val_top5': top5_m.avg, 'epoch': epoch})

## log the scaling factor of weight and activation

metrics = OrderedDict([('loss', losses_m.avg), ('top1', top1_m.avg), ('top5', top5_m.avg)])