Apex: amp.initialize

[mbaharan]

First thanks for the Apex library. It is an excellent help for FP16 training.

I am trying to load a pre-trained network again, and my concern is that if I want to use amp.initialize function, I also have to define an optimizer since it is a mandatory argument to the function. If I wish to not to use this function, I should manually cast all input to FP16 and the model itself. I am wondering if there is an API, so I can load the pre-trained network in respect to opt-level, without defining an optimizer during inference part.
Thanks.

[mcarilli]

This makes sense. I still need to implement opt-level-based checkpointing. Was the pre-trained model created by running with Amp, or without Amp?

[mbaharan]

Thanks for the reply. I trained the network with amp, and since each opt-level has a different effect on the network layer, it makes sense to also use the amp API during inference.

[mcarilli]

Do you want to train with a certain opt_level, then do inference with a different opt_level?

[mbaharan]

Not really. Same opt-level. Although I think training with O0 and inference with O3 can be considered as direct quantization without mixed training.

[hellojialee]

+1 Looking forward to simple checkpoint such that we can load the weights and do model prediction directly.

[hellojialee]

Hi @mbaharan, could you share your inference example wrapped with amp for reference. Thank you.

[mbaharan]

Hi, the following code is a portion of my project:

map_location = (lambda storage, loc: storage)
data = torch.load(check_point_file, map_location=map_location)

models = dict(data['state_dicts'][0])
dummy_input = torch.randn(10, 3, image_h_w[0], image_h_w[1], device='cuda')
model = Model(net, pretrained=False)
model.load_state_dict(models)
model.cuda()

optimizer = optim.Adam(model.parameters()) # Redundant.

model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)

features = model(dummy_input)

If I don't want to to use amp.initialize, I should cast model and inputs manually to half precision.

[hellojialee]

@mbaharan Thank you for your sharing. Did you miss model.eval()?

[mcarilli]

I recently made the "optimizers" argument to amp.initialize optional.

[mbaharan]

Thanks for the update. I will update my code based on that.

[ptrblck]

@mbaharan, @jialee93
Checkpointing just got merged into out master branch.
Checkout the README to see an example usage.