def test_quantized_imagenet(): def get_transform(): import torchvision.transforms as transforms normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) return transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), normalize, ]) def get_real_image(im_height, im_width): repo_base = 'https://github.com/dmlc/web-data/raw/master/tensorflow/models/InceptionV1/' img_name = 'elephant-299.jpg' image_url = os.path.join(repo_base, img_name) img_path = download_testdata(image_url, img_name, module='data') return Image.open(img_path).resize((im_height, im_width)) def get_imagenet_input(): im = get_real_image(224, 224) preprocess = get_transform() pt_tensor = preprocess(im) return np.expand_dims(pt_tensor.numpy(), 0) from torchvision.models.quantization import resnet as qresnet from torchvision.models.quantization import mobilenet as qmobilenet from torchvision.models.quantization import inception as qinception from torchvision.models.quantization import googlenet as qgooglenet qmodels = [] for per_channel in [False, True]: qmodels += [ ("resnet18", qresnet.resnet18(pretrained=True), per_channel), ("mobilenet_v2", qmobilenet.mobilenet_v2(pretrained=True), per_channel), # disable inception test for now, since loading it takes ~5min on torchvision-0.5 #("inception_v3", qinception.inception_v3(pretrained=True), per_channel), ("googlenet", qgooglenet(pretrained=True), per_channel), ] results = [] for (model_name, raw_model, per_channel) in qmodels: raw_model.eval() if per_channel: model_name += ", per channel quantization" else: model_name += ", per tensor quantization" inp = get_imagenet_input() pt_inp = torch.from_numpy(inp) quantize_model(raw_model, pt_inp, per_channel=per_channel, dummy=False) script_module = torch.jit.trace(raw_model, pt_inp).eval() with torch.no_grad(): pt_result = script_module(pt_inp).numpy() input_name = get_graph_input_names(script_module)[0] runtime = get_tvm_runtime(script_module, input_name, (1, 3, 224, 224)) runtime.set_input(input_name, inp) runtime.run() tvm_result = runtime.get_output(0).asnumpy() results.append((model_name, pt_result[0], tvm_result[0])) for (model_name, pt_result, tvm_result) in results: max_abs_diff = np.max(np.abs(tvm_result - pt_result)) mean_abs_diff = np.mean(np.abs(tvm_result - pt_result)) num_identical = np.sum(tvm_result == pt_result) pt_top3_labels = np.argsort(pt_result)[::-1][:3] tvm_top3_labels = np.argsort(pt_result)[::-1][:3] print("\nModel name: %s" % model_name) print("PyTorch top3 label:", pt_top3_labels) print("TVM top3 label:", tvm_top3_labels) print("max abs diff:", max_abs_diff) print("mean abs_diff:", mean_abs_diff) print("%d in 1000 raw outputs identical." % num_identical) assert set(pt_top3_labels) == set(tvm_top3_labels) # sample outputs """
pt_model(inp) t2 = time.time() print("Torch elapsed ms:", (t2 - t1) * 1e3 / n_repeat) msg = """ Loading inception v3 models on torch 1.4 + torchvision 0.5 takes a very long time (~5min). Remove "inception_v3" below to speed up testing. """ logging.warning(msg) # Mobilenet v2 was trained using QAT, post training calibration is disabled qmodels = [ ("resnet18", False, qresnet.resnet18(pretrained=True).eval()), ("resnet50", False, qresnet.resnet50(pretrained=True).eval()), ("mobilenet_v2", True, qmobilenet.mobilenet_v2(pretrained=True).eval()), ("inception_v3", False, qinception.inception_v3(pretrained=True).eval()), ("googlenet", False, qgooglenet(pretrained=True).eval()), ] if torch_version_check(): print("Adding Mobilenet v3 test") import sys sys.path.append("../models") from qmobilenet_v3 import load_model model_file = "../data/mobilenetv3small-f3be529c.pth" qmodels.append( ("mobilenet_v3 small", False, load_model(model_file).eval())) else: print("Mobilenet v3 test requires a nightly build via pip, skipping.") do_perf_bench = False