Function train

examples/pytorch/vit/ViT-quantization/main.py:196–325

(args, config)

    logger.info(f'Model is saved to {output_model_path}')

def train(args, config):
    num_classes = 1000
    model_config = CONFIGS[args.model_type]
    model = VisionTransformerINT8(model_config, args.img_size, zero_head=False, num_classes=num_classes)
    model_ckpt = torch.load(args.pretrained_dir, map_location="cpu")
    model.load_state_dict(model_ckpt["model"] if "model" in model_ckpt else model_ckpt, strict=False)
    model.cuda()
    model.train()
    
    teacher = None
    dis_loss = None
    if args.distill:
        teacher = VisionTransformerINT8(model_config, args.img_size, zero_head=False, num_classes=num_classes)
        teacher.load_from(np.load(args.teacher))
        dis_loss = Knowledge_Distillation_Loss(scale=args.distillation_loss_scale).cuda()
        teacher.cuda()
        teacher.eval()
        quant_utils.set_quantizer_by_name(teacher, [''], _disabled=True)

    """ Train the model """
    if args.local_rank in [-1, 0]:
        os.makedirs(args.output_dir, exist_ok=True)

    args.train_batch_size = args.train_batch_size // args.gradient_accumulation_steps

    # Prepare dataset
    # train_loader, test_loader = get_loader(args)
    dataset_train, dataset_val, train_loader, test_loader = build_loader(config, args)

    # Prepare optimizer and scheduler
    optimizer = torch.optim.SGD(model.parameters(),
                                lr=args.qat_lr,
                                momentum=0.9,
                                weight_decay=args.weight_decay)
    
    print('args.qat_lr: %.6f' % (args.qat_lr))
    print('optimizer.lr: %.6f' % optimizer.state_dict()['param_groups'][0]['lr'])
    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)

    
    print('查看optimizer.param_groups结构:')
    i_list=[i for i in optimizer.param_groups[0].keys()]
    print(i_list)
    if args.fp16:
        model, optimizer = amp.initialize(models=model,
                                          optimizers=optimizer,
                                          opt_level=args.fp16_opt_level)
        amp._amp_state.loss_scalers[0]._loss_scale = 2**20

    # Distributed training
    if args.local_rank != -1:
        model = DDP(model, message_size=250000000, gradient_predivide_factor=get_world_size())

    # Train!