# 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 = {

# "enable": args.wq_enable,

"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, boundaryRange= args.boundaryRange)

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, boundaryRange= args.boundaryRange)

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)

not_freeze_idx = torch.zeros_like(cur_weight).to(cur_weight.device)

clip_val = torch.tensor([2.0]).to(cur_weight.device)

def get_min_max_round(weights, scaling_factor, clip_val, num_bits):

scaled_weights = weights/scaling_factor

cliped_weights = torch.clamp(scaled_weights, min=(-clip_val/2), max=(clip_val/2)-1e-6)

n = float(2 ** (num_bits - 1))

b4_round = (cliped_weights) * n - 0.5

round_result = torch.round(b4_round)

return b4_round, round_result, torch.min(round_result), torch.max(round_result)

scaling_factor = 2 * torch.mean(abs(cur_weight),dim=1,keepdim=True) # dim, 1

b4_round, round, min, max = get_min_max_round(weights=cur_weight, scaling_factor=scaling_factor, clip_val=clip_val,num_bits=num_bits)

for i in np.arange(start=min.detach().cpu(),stop=max.detach().cpu()): # 0.5 - boundaryRange < x < 0.5 + boundaryRange

within_boundary = ((b4_round - i) <= (0.5 + boundaryRange)) * ((b4_round - i) >= (0.5 - boundaryRange)) #idx of # 0.5 - boundaryRange < x < 0.5 + boundaryRange

not_freeze_idx = not_freeze_idx + within_boundary.float()

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

args.cooldown_epochs = args.freeze_for_n_epochs

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

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 = args.epochs

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)

# setup checkpoint saver and eval metric tracking

eval_metric = args.eval_metric

best_metric = None

best_epoch = None

saver = None

output_dir = None

if args.rank == 0:

if args.experiment:

exp_name = args.experiment

else:

exp_name = '-'.join([

datetime.now().strftime("%Y%m%d-%H%M%S"),

safe_model_name(args.model),

str(data_config['input_size'][-1])

])

output_dir = get_outdir(args.output if args.output else './output/train', exp_name)

decreasing = True if eval_metric == 'loss' else False

saver = CheckpointSaver(

model=model, optimizer=optimizer, args=args, model_ema=model_ema, amp_scaler=loss_scaler,

checkpoint_dir=output_dir, recovery_dir=output_dir, decreasing=decreasing, max_history=args.checkpoint_hist)

with open(os.path.join(output_dir, 'args.yaml'), 'w') as f:

f.write(args_text)

try:

if args.local_rank == 0:

print('start_epoch: ', start_epoch)

num_epochs = start_epoch + args.freeze_for_n_epochs

print('total epoch: ',num_epochs)

for epoch in range(start_epoch, num_epochs):

if args.distributed and hasattr(loader_train.sampler, 'set_epoch'):

loader_train.sampler.set_epoch(epoch)

train_metrics = train_one_epoch(

epoch, model, loader_train, optimizer, train_loss_fn, args,

lr_scheduler=lr_scheduler, saver=saver, output_dir=output_dir,

amp_autocast=amp_autocast, loss_scaler=loss_scaler, model_ema=model_ema,

mixup_fn=mixup_fn, teacher=teacher,teacher_type=args.teacher_type, log_wandb=args.log_wandb and has_wandb )

if args.distributed and args.dist_bn in ('broadcast', 'reduce'):

if args.local_rank == 0:

_logger.info("Distributing BatchNorm running means and vars")

distribute_bn(model, args.world_size, args.dist_bn == 'reduce')

eval_metrics = validate(model, loader_eval, validate_loss_fn, args, amp_autocast=amp_autocast,epoch=epoch, log_wandb=args.log_wandb and has_wandb)

if model_ema is not None and not args.model_ema_force_cpu:

if args.distributed and args.dist_bn in ('broadcast', 'reduce'):

distribute_bn(model_ema, args.world_size, args.dist_bn == 'reduce')

ema_eval_metrics = validate(

model_ema.module, loader_eval, validate_loss_fn, args, amp_autocast=amp_autocast,

log_suffix=' (EMA)')

eval_metrics = ema_eval_metrics

if lr_scheduler is not None:

# step LR for next epoch

lr_scheduler.step(epoch + 1, eval_metrics[eval_metric])

if output_dir is not None:

update_summary(

epoch, train_metrics, eval_metrics, os.path.join(output_dir, 'summary.csv'),

write_header=best_metric is None, log_wandb=args.log_wandb and has_wandb)

if saver is not None:

# save proper checkpoint with eval metric

save_metric = eval_metrics[eval_metric]

best_metric, best_epoch = saver.save_checkpoint(epoch, metric=save_metric)

except KeyboardInterrupt:

pass

if best_metric is not None:

_logger.info('*** Best metric: {0} (epoch {1})'.format(best_metric, best_epoch))

if args.log_wandb and has_wandb:

epoch, model, loader, optimizer, loss_fn, args,

lr_scheduler=None, saver=None, output_dir=None, amp_autocast=suppress,

loss_scaler=None, model_ema=None, mixup_fn=None, teacher=None, teacher_type = 'deit', log_wandb = False, annealing_schedule_for_dampen = None

, annealing_schedule_for_iterative_freezing = None):

if args.mixup_off_epoch and epoch >= args.mixup_off_epoch:

if args.prefetcher and loader.mixup_enabled:

loader.mixup_enabled = False

elif mixup_fn is not None:

mixup_fn.mixup_enabled = False

second_order = hasattr(optimizer, 'is_second_order') and optimizer.is_second_order

batch_time_m = AverageMeter()

data_time_m = AverageMeter()

losses_m = AverageMeter()

model.train()

end = time.time()

last_idx = len(loader) - 1

num_updates = epoch * len(loader)

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

last_batch = batch_idx == last_idx

data_time_m.update(time.time() - end)

if not args.prefetcher:

input, target = input.cuda(), target.cuda()

if mixup_fn is not None:

input, target = mixup_fn(input, target)

if args.channels_last:

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

with amp_autocast():

if args.model_type == 'deit' or args.model_type == 'swin':

student_logit, student_attn_info = model(input)

else:

student_logit = model(input)

if args.use_kd:

if args.kd_hard_and_soft == 0:

if args.teacher_type == 'deit':

soft_target, _ = teacher(input)

else:

soft_target = teacher(input)

loss = loss_fn(student_logit, soft_target)

elif args.kd_hard_and_soft == 1:

if args.teacher_type == 'deit' or args.teacher_type == 'swin':

soft_target, teacher_attn_info = teacher(input)

else:

soft_target = teacher(input)

loss = loss_fn(student_logit, target, soft_target)

else:

student_logit = student_logit[0] if isinstance(student_logit, tuple) else student_logit

loss = loss_fn(student_logit, target)

if not args.distributed:

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

optimizer.zero_grad()

if loss_scaler is not None:

loss_scaler(

loss, optimizer,

clip_grad=args.clip_grad, clip_mode=args.clip_mode,

parameters=model_parameters(model, exclude_head='agc' in args.clip_mode),

create_graph=second_order)

else:

loss.backward(create_graph=second_order)

### set the gradient of the weight outside the boundary to be 0 args.boundaryRange

save_frozen_weight = {}

freeze_idx_dic = {}

if args.qk_reparam:

if args.model_type == 'swin':

for k, v in model.named_modules():

if (k[-3:] == 'fc1' or k[-3:] == 'fc2' or k[-2:] == '.v' or k[-4:] == 'proj' or k[-9:] == 'reduction'):

freeze_idx = freeze_outside_boundary_weight_idx( v.weight, args.wq_bitw, boundaryRange = args.boundaryRange)

freeze_idx_dic[k]= freeze_idx.detach().clone()

v.weight.grad = v.weight.grad * freeze_idx * 0.0 + v.weight.grad * (1 - freeze_idx )

save_frozen_weight[k] = (v.weight * freeze_idx).detach().clone()

else:

for k, v in model.named_modules():

if 'blocks' in k and (k[-3:] == 'fc1' or k[-3:] == 'fc2' or k[-2:] == '.v' or k[-4:] == 'proj'):

freeze_idx = freeze_outside_boundary_weight_idx( v.weight, args.wq_bitw, boundaryRange = args.boundaryRange)

freeze_idx_dic[k]= freeze_idx.detach().clone()

v.weight.grad = v.weight.grad * freeze_idx * 0.0 + v.weight.grad * (1 - freeze_idx )

save_frozen_weight[k] = (v.weight * freeze_idx).detach().clone()

else:

for k, v in model.named_modules():

if 'blocks' in k and (k[-3:] == 'fc1' or k[-3:] == 'fc2' or k[-3:] == 'qkv' or k[-4:] == 'proj'):

freeze_idx = freeze_outside_boundary_weight_idx( v.weight, args.wq_bitw, boundaryRange = args.boundaryRange)

freeze_idx_dic[k]= freeze_idx.detach().clone()

v.weight.grad = v.weight.grad * freeze_idx * 0.0 + v.weight.grad * (1 - freeze_idx )

save_frozen_weight[k] = (v.weight * freeze_idx).detach().clone()

if args.clip_grad is not None:

dispatch_clip_grad(

model_parameters(model, exclude_head='agc' in args.clip_mode),

value=args.clip_grad, mode=args.clip_mode)

optimizer.step()

## restore the forzen weight:

if args.qk_reparam:

if args.model_type == 'swin':

for k, v in model.named_modules():

if (k[-3:] == 'fc1' or k[-3:] == 'fc2' or k[-2:] == '.v' or k[-4:] == 'proj' or k[-9:] == 'reduction'):

with torch.no_grad():

keep_w = (v.weight.detach().clone() * (1-freeze_idx_dic[k]))

replace_freeze_part = save_frozen_weight[k]

new_weight = keep_w + replace_freeze_part

v.weight.data.copy_(new_weight)

else:

for k, v in model.named_modules():

if 'blocks' in k and (k[-3:] == 'fc1' or k[-3:] == 'fc2' or k[-2:] == '.v' or k[-4:] == 'proj'):