Fast histogram observer
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/29790
Test Plan:
import torch
import time
import numpy as np
from torch.quantization.observer import HistogramObserver
X = torch.randn(1,1,224,224)
obs = HistogramObserver(2048)
acc_time = 0
for i in range(100):
X = torch.randn(10,1,320,320)
start = time.time()
obs(X)
#obs.forward_new(X)
acc_time = acc_time + time.time()-start
print(acc_time)
Imported from OSS
Differential Revision: D18508562
fbshipit-source-id: 456e82360ce1b3f9d8b6e1832d23f1339655011a