When i follow tutorial from https://github.com/pytorch/vision/blob/39021408587eb252ebb842a54195d840d6d76095/references/classification/train_quantization.py, when i reload pretrained qat model, it will cause an error:
Traceback (most recent call last):
File "train_quantization.py", line 461, in <module>
main(args)
File "train_quantization.py", line 311, in main
args.print_freq, priors)
File "train_quantization.py", line 67, in train_one_epoch
optimizer.step()
File "/opt/conda/lib/python3.7/site-packages/torch/optim/lr_scheduler.py", line 67, in wrapper
return wrapped(*args, **kwargs)
File "/opt/conda/lib/python3.7/site-packages/torch/autograd/grad_mode.py", line 15, in decorate_context
return func(*args, **kwargs)
File "/opt/conda/lib/python3.7/site-packages/torch/optim/sgd.py", line 106, in step
buf.mul_(momentum).add_(d_p, alpha=1 - dampening)
RuntimeError: expected device cpu but got device cuda:0
I notice the optimizer reload error, when i use a new optimizer it will be ok
if args.resume:
checkpoint = torch.load(args.resume, map_location='cpu')
model_without_ddp.load_state_dict(checkpoint['model'])
#optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
args.start_epoch = checkpoint['epoch'] + 1
optimizer = torch.optim.SGD(
model.parameters(), lr=args.lr, momentum=args.momentum,
weight_decay=args.weight_decay)
model.to(device)
I use the newest pytorch docker image form dockerhub.