Calibrate the model
(args, config, model)
| 165 | return acc1_meter.avg |
| 166 | |
| 167 | def calib(args, config, model): |
| 168 | """ Calibrate the model """ |
| 169 | if args.local_rank in [-1, 0]: |
| 170 | os.makedirs(args.output_dir, exist_ok=True) |
| 171 | |
| 172 | dataset_train, dataset_val, train_loader, test_loader = build_loader(config, args) |
| 173 | # Calibration |
| 174 | quant_utils.configure_model(model, args, calib=True) |
| 175 | model.eval() |
| 176 | quant_utils.enable_calibration(model) |
| 177 | # Run forward passes on a sample of the training set |
| 178 | for step, (samples, targets) in enumerate(tqdm(train_loader, desc='Calibration', total=args.num_calib_batch)): |
| 179 | if step > args.num_calib_batch: |
| 180 | break |
| 181 | samples = samples.to(args.device) |
| 182 | outputs = model(samples) |
| 183 | quant_utils.finish_calibration(model, args) |
| 184 | |
| 185 | # model.load_state_dict(torch.load('checkpoint/{}_{}_{}.pth'.format(args.model_type, args.quant_mode, args.percentile))) |
| 186 | |
| 187 | quant_utils.configure_model(model, args, calib=False) |
| 188 | if args.local_rank in [-1, 0]: |
| 189 | accuracy = valid(args, config, model, test_loader) |
| 190 | logger.info("Test Accuracy: \t%f" %accuracy) |
| 191 | |
| 192 | output_model_path = os.path.join(args.calib_output_path, '{}_calib.pth'.format(args.model_type)) |
| 193 | if not os.path.exists(args.calib_output_path): |
| 194 | os.mkdir(args.calib_output_path) |
| 195 | torch.save(model.state_dict(), output_model_path) |
| 196 | logger.info(f'Model is saved to {output_model_path}') |
| 197 | |
| 198 | def validate_trt(args, config): |
| 199 | num_classes = 1000 |
no test coverage detected