Fuse Conv+BN layers in pretrained model and save it as new model


I would like to load a pretrained model, fuse the Conv+BN layers, and save it as a new model.

I have found functions that perform the merging but it seems that I miss the workflow which is pytorch specific, e.g.

  1. add fusion in the model constructor
  2. load_state_dict → complains with error because it expects different layer configuration (not fused)

It seems that I should also have modify the model.state_dict but I could not find how this is done.

Thanks in advance,

optimize_for_inference() does this for you torch.jit.optimize_for_inference — PyTorch 1.10 documentation