def verify_script_model(pt_model, ishapes):
script_module = torch.jit.script(pt_model)
input_names = get_graph_input_names(script_module)
input_shapes = dict(zip(input_names, ishapes))
inputs = [
torch.randn(input_shapes[input_name], dtype=torch.float)
for input_name in input_names
]
mod, params = relay.frontend.from_pytorch(script_module, input_shapes)
executor = relay.create_executor("vm",
mod=mod,
ctx=tvm.cpu(0),
target="llvm")
evaluator = executor.evaluate()
for name, inp in zip(input_names, inputs):
params[name] = inp.numpy()
op_res = evaluator(**params)
with torch.no_grad():
pt_result = pt_model(*inputs)
if not isinstance(pt_result, torch.Tensor):
tvm_res = op_res.asnumpy().item()
assert pt_result == tvm_res
else:
tvm.testing.assert_allclose(op_res.asnumpy(),
pt_result.numpy(),
rtol=1e-5,
atol=1e-5)
def run_on_models(models, inputs, target="llvm"):
for raw_model in models:
with torch.no_grad():
pt_result = raw_model(*inputs).numpy()
script_module = torch.jit.trace(raw_model, *inputs).eval()
input_names = get_graph_input_names(script_module)
input_shapes = dict(zip(input_names, [inp.shape for inp in inputs]))
mod, params = from_pytorch(script_module, input_shapes)
with relay.build_config(opt_level=3):
json, lib, params = relay.build(mod, target=target, params=params)
ctx = tvm.context(target, 0)
runtime = tvm.contrib.graph_runtime.create(json, lib, ctx)
runtime.set_input(**params)
for name, inp in zip(input_names, inputs):
runtime.set_input(name, inp.numpy())
runtime.run()
tvm_result = runtime.get_output(0).asnumpy()
np.allclose(tvm_result, pt_result, rtol=1e-5, atol=1e-5)
print(np.max(np.abs(tvm_result - pt_result)),
np.mean(np.abs(tvm_result - pt_result)))
tvm.testing.assert_allclose(tvm_result, pt_result,
rtol=1e-3, atol=1e-3)
def test_quantized_modules():
imagenet_ishape = (1, 3, 224, 224)
qmodules = [
("relu", imagenet_ishape, ReLU(), False),
("upsample bilinear", (1, 3, 64, 64), UpsamplingBilinear(), False),
]
for per_channel in [False, True]:
if per_channel:
postfix = ", per_channel"
else:
postfix = ""
qmodules += [
("conv_bn" + postfix, imagenet_ishape, ConvBn(), per_channel),
("conv_bn_relu" + postfix, imagenet_ishape, ConvBn(with_relu=True), per_channel),
("linear" + postfix, (16, 16), Linear(), per_channel),
("linear_relu" + postfix, (16, 16), Linear(with_relu=True), per_channel)
]
if torch_version_check():
qmodules += [
("hsigmoid", imagenet_ishape, Hsigmoid(add_stub=True), False),
("hswish", imagenet_ishape, Hswish(add_stub=True), False),
("semodule", (1, 16, 64, 64), SqueezeExcite(16, add_stub=True), False),
("semodule, per_channel", (1, 16, 64, 64), SqueezeExcite(16, add_stub=True), True),
("mul_scalar negative", imagenet_ishape, MulScalarNegative(), False)
]
else:
print("Skipping tests that require torch > 1.4")
for (module_name, ishape, raw_module, per_channel) in qmodules:
raw_module.eval()
inp = torch.rand(ishape)
quantize_model(raw_module, inp, per_channel=per_channel, dummy=True)
script_module = torch.jit.trace(raw_module, inp).eval()
with torch.no_grad():
pt_result = script_module(inp.clone()).numpy()
input_name = get_graph_input_names(script_module)[0]
runtime = get_tvm_runtime(script_module, input_name, ishape)
runtime.set_input(input_name, inp.numpy().copy())
runtime.run()
tvm_result = runtime.get_output(0).asnumpy()
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)
match_ratio = num_identical / float(np.prod(tvm_result.shape))
print(module_name, max_abs_diff, mean_abs_diff, match_ratio)
# sample outputs
"""
def verify_model(model_name, input_data=[]):
"""Assert that the output of a compiled model matches with that of its
baseline."""
if len(input_data) == 0:
baseline_model, baseline_input = load_model(model_name)
else:
baseline_model = model_name
baseline_input = input_data
if torch.cuda.is_available():
baseline_model = baseline_model.cuda()
baseline_input = baseline_input.cuda()
with torch.no_grad():
baseline_outputs = baseline_model(baseline_input)
if isinstance(baseline_outputs, tuple):
baseline_outputs = tuple(out.cpu().numpy() for out in baseline_outputs)
else:
baseline_outputs = (baseline_outputs.float().cpu().numpy(),)
trace = torch.jit.trace(baseline_model, baseline_input).float().eval()
if torch.cuda.is_available():
trace = trace.cuda()
else:
trace = trace.cpu()
input_name = get_graph_input_names(trace)[0] # only one input
input_shapes = {input_name: list(baseline_input.shape)}
mod, params = relay.frontend.from_pytorch(trace, input_shapes)
compiled_input = {input_name: tvm.nd.array(baseline_input.cpu().numpy())}
with relay.build_config(opt_level=3):
for target, ctx in ctx_list():
relay_graph, relay_lib, relay_params = relay.build(mod, target=target, params=params)
relay_model = graph_runtime.create(relay_graph, relay_lib, ctx)
relay_model.set_input(**relay_params)
relay_model.set_input(**compiled_input)
relay_model.run()
for i, baseline_output in enumerate(baseline_outputs):
compiled_output = relay_model.get_output(i).asnumpy()
assert_shapes_match(baseline_output, compiled_output)
tvm.testing.assert_allclose(baseline_output, compiled_output,
rtol=1e-3, atol=1e-3)
del model_name
del baseline_model
torch.cuda.empty_cache()
def simple_rnn_test():
class DecisionGate(torch.nn.Module):
def forward(self, x):
if x.sum() > 0:
return x
else:
return -x
class Cell(torch.nn.Module):
def __init__(self, dg):
super(Cell, self).__init__()
self.dg = dg
self.linear = torch.nn.Linear(4, 4)
def forward(self, x, h):
new_h = torch.tanh(self.dg(self.linear(x)) + h)
return new_h, new_h
class RNNLoop(torch.nn.Module):
def __init__(self):
super().__init__()
x = torch.rand(10, 4, dtype=torch.float)
h = torch.rand(10, 4, dtype=torch.float)
self.cell = torch.jit.trace(Cell(DecisionGate()), (x, h))
def forward(self, xs):
h = torch.zeros(10, 4, dtype=torch.float)
y = torch.zeros(10, 4, dtype=torch.float)
for i in range(xs.size(0)):
y, h = self.cell(xs[i], h)
return y
raw_model = RNNLoop()
script_module = torch.jit.script(raw_model)
input_name = get_graph_input_names(script_module)[0]
input_shapes = {input_name: (10, 10, 4)}
mod, params = from_pytorch(script_module, input_shapes, {})
for i in range(5):
inp = torch.randn(input_shapes[input_name], dtype=torch.float)
with torch.no_grad():
pt_result = raw_model(inp.clone())
params[input_name] = inp.numpy()
run_and_compare(mod, params, pt_result)
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
"""
# Preprocess the image and convert to tensor
from torchvision import transforms
my_preprocess = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
img = my_preprocess(img)
img = np.expand_dims(img, 0)
######################################################################
# Import the graph to Relay
# -------------------------
# Convert PyTorch graph to Relay graph.
input_name = get_graph_input_names(scripted_model)[0] # only one input
shape_dict = {input_name: img.shape}
mod, params = relay.frontend.from_pytorch(scripted_model, shape_dict)
######################################################################
# Relay Build
# -----------
# Compile the graph to llvm target with given input specification.
target = 'llvm'
target_host = 'llvm'
ctx = tvm.cpu(0)
with relay.build_config(opt_level=3):
graph, lib, params = relay.build(mod,
target=target,
target_host=target_host,
params=params)
def verify_model(model_name,
input_data=[],
custom_convert_map={},
ctx_list=ctx_list()):
"""Assert that the output of a compiled model matches with that of its
baseline."""
if isinstance(model_name, str):
baseline_model, baseline_input = load_model(model_name)
elif isinstance(input_data, list):
baseline_model = model_name
baseline_input = input_data
elif isinstance(input_data, torch.Tensor) or len(input_data.shape) == 0:
baseline_model = model_name
baseline_input = [input_data]
else:
assert False, "Unexpected input format"
if torch.cuda.is_available():
baseline_model = baseline_model.cuda()
baseline_input = [inp.cuda() for inp in baseline_input]
with torch.no_grad():
baseline_outputs = baseline_model(*baseline_input)
if isinstance(baseline_outputs, tuple):
baseline_outputs = tuple(out.cpu().numpy() for out in baseline_outputs)
else:
baseline_outputs = (baseline_outputs.float().cpu().numpy(), )
trace = torch.jit.trace(baseline_model, baseline_input).float().eval()
if torch.cuda.is_available():
trace = trace.cuda()
else:
trace = trace.cpu()
input_names = get_graph_input_names(trace)
input_shapes = dict(zip(input_names,
[inp.shape for inp in baseline_input]))
mod, params = relay.frontend.from_pytorch(trace, input_shapes,
custom_convert_map)
compiled_input = dict(
zip(input_names, [inp.cpu().numpy() for inp in baseline_input]))
with relay.build_config(opt_level=3):
for target, ctx in ctx_list:
relay_graph, relay_lib, relay_params = relay.build(mod,
target=target,
params=params)
relay_model = graph_runtime.create(relay_graph, relay_lib, ctx)
relay_model.set_input(**relay_params)
for name, inp in compiled_input.items():
relay_model.set_input(name, inp)
relay_model.run()
for i, baseline_output in enumerate(baseline_outputs):
compiled_output = relay_model.get_output(i).asnumpy()
assert_shapes_match(baseline_output, compiled_output)
tvm.testing.assert_allclose(baseline_output,
compiled_output,
rtol=1e-3,
atol=1e-3)
del model_name
del baseline_model
torch.cuda.empty_cache()
models = [
SimpleIf(10, 20).eval(),
NestedIf(10, 20).eval(),
ScalarLoop().eval(),
SimpleLoop().eval(),
LoopWithIf().eval(),
SimpleScalarWhileLoop().eval(),
SimpleWhileLoop().eval(),
NestedLoop().eval()
]
for raw_model in models:
script_module = torch.jit.script(raw_model)
input_name = get_graph_input_names(script_module)[0]
input_shapes = {input_name: (10, 20)}
mod, params = from_pytorch(script_module, input_shapes)
print(mod["main"])
executor = relay.create_executor("vm", mod=mod, ctx=tvm.cpu(0), target="llvm")
evaluator = executor.evaluate()
for i in range(5):
inp = torch.rand(input_shapes[input_name], dtype=torch.float)
with torch.no_grad():
pt_result = raw_model(inp.clone())
params[input_name] = inp.numpy()
op_res = evaluator(**params)
