(args, config)
| 299 | return accuracy |
| 300 | |
| 301 | def train(args, config): |
| 302 | num_classes = 1000 |
| 303 | model_config = CONFIGS[args.model_type] |
| 304 | model = VisionTransformerINT8(model_config, args.img_size, zero_head=False, num_classes=num_classes) |
| 305 | model_ckpt = torch.load(args.pretrained_dir, map_location="cpu") |
| 306 | model.load_state_dict(model_ckpt["model"] if "model" in model_ckpt else model_ckpt, strict=False) |
| 307 | model.cuda() |
| 308 | model.train() |
| 309 | |
| 310 | teacher = None |
| 311 | dis_loss = None |
| 312 | if args.distill: |
| 313 | teacher = VisionTransformerINT8(model_config, args.img_size, zero_head=False, num_classes=num_classes) |
| 314 | teacher.load_from(np.load(args.teacher)) |
| 315 | dis_loss = Knowledge_Distillation_Loss(scale=args.distillation_loss_scale).cuda() |
| 316 | teacher.cuda() |
| 317 | teacher.eval() |
| 318 | quant_utils.set_quantizer_by_name(teacher, [''], _disabled=True) |
| 319 | |
| 320 | """ Train the model """ |
| 321 | if args.local_rank in [-1, 0]: |
| 322 | os.makedirs(args.output_dir, exist_ok=True) |
| 323 | |
| 324 | args.train_batch_size = args.train_batch_size // args.gradient_accumulation_steps |
| 325 | |
| 326 | # Prepare dataset |
| 327 | # train_loader, test_loader = get_loader(args) |
| 328 | dataset_train, dataset_val, train_loader, test_loader = build_loader(config, args) |
| 329 | |
| 330 | # Prepare optimizer and scheduler |
| 331 | optimizer = torch.optim.SGD(model.parameters(), |
| 332 | lr=args.qat_lr, |
| 333 | momentum=0.9, |
| 334 | weight_decay=args.weight_decay) |
| 335 | |
| 336 | print('args.qat_lr: %.6f' % (args.qat_lr)) |
| 337 | print('optimizer.lr: %.6f' % optimizer.state_dict()['param_groups'][0]['lr']) |
| 338 | t_total = args.num_steps |
| 339 | # if args.decay_type == "cosine": |
| 340 | # scheduler = WarmupCosineSchedule(optimizer, warmup_steps=args.warmup_steps, t_total=t_total) |
| 341 | # else: |
| 342 | # scheduler = WarmupLinearSchedule(optimizer, warmup_steps=args.warmup_steps, t_total=t_total) |
| 343 | |
| 344 | |
| 345 | print('查看optimizer.param_groups结构:') |
| 346 | i_list=[i for i in optimizer.param_groups[0].keys()] |
| 347 | print(i_list) |
| 348 | if args.fp16: |
| 349 | model, optimizer = amp.initialize(models=model, |
| 350 | optimizers=optimizer, |
| 351 | opt_level=args.fp16_opt_level) |
| 352 | amp._amp_state.loss_scalers[0]._loss_scale = 2**20 |
| 353 | |
| 354 | # Distributed training |
| 355 | if args.local_rank != -1: |
| 356 | model = DDP(model, message_size=250000000, gradient_predivide_factor=get_world_size()) |
| 357 | |
| 358 | # Train! |
no test coverage detected