Skip to content

Instantly share code, notes, and snippets.

@vkuzo
Created January 27, 2021 16:47
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save vkuzo/8ad0e472e8743e142a1e72cb0f9efb1d to your computer and use it in GitHub Desktop.
Save vkuzo/8ad0e472e8743e142a1e72cb0f9efb1d to your computer and use it in GitHub Desktop.
import torch
import torch.nn as nn
m = nn.Sequential(nn.Conv2d(1, 1, 1))
m.qconfig = torch.quantization.default_qconfig
mp = torch.quantization.prepare(m)
mp(torch.randn(4, 1, 4, 4))
mq = torch.quantization.convert(mp)
print('before', mq[0].state_dict())
w, b = mq[0]._weight_bias()
new_w = torch.quantize_per_tensor(torch.randn(1, 1, 1, 1), 0.1, 0, torch.qint8)
mq[0].set_weight_bias(new_w, b)
print('after', mq[0].state_dict())
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment