def tuple_test_helper(self, ModType): input = torch.randn(4) model = ModType() spec = torch_glow.CompilationSpec() spec.get_settings().set_glow_backend("Interpreter") compilation_group = torch_glow.CompilationGroup() spec.compilation_groups_append(compilation_group) input_spec = torch_glow.InputSpec() input_spec.set_same_as(input) compilation_group.input_sets_append([input_spec]) scripted_mod = torch.jit.script(model) lowered_model = torch_glow.to_glow(scripted_mod, {"forward": spec}) # Run Glow model g = lowered_model(input) # Run reference model t = model(input) self.assertEqual(type(g), type(t)) self.assertEqual(len(g), len(t)) for (gi, ti) in zip(g, t): self.assertTrue(torch.allclose(gi, ti))
def test_save_preprocessed_module(self): with torch.no_grad(): x = torch.randn([1, 4, 4, 4], dtype=torch.float32) model = Bar() model.eval() model = torch.jit.trace(model, x) spec = torch_glow.CompilationSpec() spec.get_settings().set_glow_backend("Interpreter") compilation_group = torch_glow.CompilationGroup() spec.compilation_groups_append(compilation_group) compilation_group.input_sets_append( torch_glow.input_specs_from_tensors([x])) torch_glow.disableFusionPass() torch_glow.enable_convert_to_fp16() glow_mod = torch_glow.to_glow(model, spec) reloaded = utils.save_and_reload_model(glow_mod) wrappername = "__loweredModule__" attrname = "__processed_module" wp = getattr(reloaded._c, wrappername) pp = getattr(wp, attrname) pt_model = torch.jit._recursive.wrap_cpp_module(pp) graph = pt_model.graph_for(x) found = False for node in graph.nodes(): if node.kind() == "quantized::conv2d": found = True assert found
def test_to_glow_tuple_output(self): input = torch.randn(4) model = Foo() spec = torch_glow.CompilationSpec() spec.get_settings().set_glow_backend("Interpreter") compilation_group = torch_glow.CompilationGroup() spec.compilation_groups_append(compilation_group) input_spec = torch_glow.InputSpec() input_spec.set_same_as(input) compilation_group.input_sets_append([input_spec]) scripted_mod = torch.jit.script(model) lowered_model = torch_glow.to_glow(scripted_mod, {"forward": spec}) # Run Glow model (gx, gy) = lowered_model(input) # Run reference model (tx, ty) = model(input) assert torch.allclose(tx, gx) assert torch.allclose(ty, gy)
def devices_to_use_test_helper(self, input, num_replications): model = SimpleModule() spec = torch_glow.CompilationSpec() spec.get_settings().set_glow_backend("Interpreter") # Init with total number of devices. torch_glow.setGlowBackendNumDevices(6) compilation_group = torch_glow.CompilationGroup() spec.compilation_groups_append(compilation_group) input_spec = torch_glow.InputSpec() input_spec.set_same_as(input) compilation_group.input_sets_append([input_spec]) compilation_group_settings = compilation_group.get_settings() compilation_group_settings.set_num_devices_to_use(3) compilation_group_settings.set_replication_count(num_replications) traced_mod = torch.jit.trace(model, input) lowered_model = torch_glow.to_glow(traced_mod, {"forward": spec}) g = lowered_model(input) t = model(input) self.assertEqual(type(g), type(t)) self.assertEqual(len(g), len(t)) for (gi, ti) in zip(g, t): self.assertTrue(torch.allclose(gi, ti))
def test_to_glow_multiple_specs(self): a = torch.randn(4) b = torch.randn(6) model = Foo() torch_resA = model(a) torch_resB = model(b) metaA = torch_glow.InputMeta() metaA.set_same_as(a) inputA = [metaA] metaB = torch_glow.InputMeta() metaB.set_same_as(b) inputB = [metaB] options = torch_glow.CompilationOptions() options.backend = "Interpreter" specA = torch_glow.GlowCompileSpec() specA.set(inputA, options) specB = torch_glow.GlowCompileSpec() specB.set(inputB, options) scripted_mod = torch.jit.script(model) lowered_mod = torch_glow.to_glow(scripted_mod, [specA, specB]) glow_resA = lowered_mod(a) glow_resB = lowered_mod(b) assert torch.allclose(torch_resA, glow_resA) assert torch.allclose(torch_resB, glow_resB)
def compare_tracing_methods( module, *inputs, atol=5e-4, rtol=1e-3, reference=None, fusible_ops=None, fusion_blocklist=None, fp16=False, scripted=False, check_trace=True, skip_to_glow=False, # Ugly hack, TODO: Remove ): if not isinstance(module, torch.nn.Module): raise AssertionError("to_glow only supports nn.Modules") def trace(mod, ins): if scripted: return torch.jit.script(mod) else: return torch.jit.trace(mod, ins, check_trace=check_trace) with torch.no_grad(): with ephemeral_torchglow_settings( fusion=True, fp16=fp16, blocklist=fusion_blocklist ): fusion_inputs = deepcopy(inputs) fusion_trace = trace(module, fusion_inputs) assert_fused( fusion_trace.graph_for(*fusion_inputs), *(fusible_ops or []), accept_any=fusible_ops is None, ) fusion_result = fusion_trace(*fusion_inputs) with ephemeral_torchglow_settings(fusion=False, fp16=fp16): if scripted: torchscript_result = module(*deepcopy(inputs)) else: torchscript_inputs = deepcopy(inputs) torchscript_trace = trace(module, torchscript_inputs) torchscript_result = torchscript_trace(*torchscript_inputs) with ephemeral_torchglow_settings(fusion=False, fp16=fp16): if not skip_to_glow: glow_inputs = deepcopy(inputs) glow_spec = generate_glow_spec(module, DEFAULT_BACKEND, *glow_inputs) glow_trace = torch_glow.to_glow(trace(module, glow_inputs), glow_spec) glow_result = glow_trace(*glow_inputs) if reference: assert_equivalent(reference, fusion_trace, atol=atol, rtol=rtol) assert_equivalent(reference, torchscript_result, atol=atol, rtol=rtol) if not skip_to_glow: assert_equivalent(reference, glow_result, atol=atol, rtol=rtol) # This is written out manually instead of using combinations in order to aid # debugging. TODO: Clean up. assert_equivalent(fusion_result, torchscript_result, atol=atol, rtol=rtol) if not skip_to_glow: assert_equivalent(fusion_result, glow_result, atol=atol, rtol=rtol) assert_equivalent(torchscript_result, glow_result, atol=atol, rtol=rtol)
def compare_tracing_methods_error( module, *inputs, fusible_ops=None, fusion_blocklist=None, fp16=False, ): if not isinstance(module, torch.nn.Module): raise AssertionError("to_glow only supports nn.Modules") def trace(mod, ins): return torch.jit.trace(mod, ins) with torch.no_grad(): with ephemeral_torchglow_settings(fusion=True, fp16=fp16, blocklist=fusion_blocklist): fusion_inputs = deepcopy(inputs) try: fusion_trace = trace(module, fusion_inputs) assert_fused( fusion_trace.graph_for(*fusion_inputs), *(fusible_ops or []), accept_any=fusible_ops is None, ) fusion_trace(*fusion_inputs) except Exception: pass else: raise AssertionError( "Error expected (fusion), but none were received") with ephemeral_torchglow_settings(fusion=False, fp16=fp16): try: torchscript_inputs = deepcopy(inputs) torchscript_trace = trace(module, torchscript_inputs) torchscript_trace(*torchscript_inputs) except Exception: pass else: raise AssertionError( "Error expected (torchscript), but none were received") with ephemeral_torchglow_settings(fusion=False, fp16=fp16): try: glow_inputs = deepcopy(inputs) glow_spec = torch_glow.lower( model=module, example_inputs=glow_inputs, backend=DEFAULT_BACKEND, ) glow_trace = torch_glow.to_glow(trace(module, glow_inputs), glow_spec) glow_trace(*glow_inputs) except Exception: pass else: raise AssertionError( "Error expected (glow), but none were received")
def run_to_glow(m, x): """Trace the model m with input x and call to_glow""" traced_m = torch.jit.trace(m, (x)) spec = torch.classes.glow.GlowCompileSpec() spec.setBackend("Interpreter") sim = torch.classes.glow.SpecInputMeta() sim.set(x.size(), torch.float32) inputs = [sim] spec.addInputs(inputs) lowered_module = torch_glow.to_glow(traced_m, {"forward": spec}) return lowered_module
def run_to_glow(m, x): """Trace the model m with input x and call to_glow""" traced_m = torch.jit.trace(m, (x)) input_meta = InputMeta() input_meta.set(x.size(), torch.float32) inputs = [input_meta] options = CompilationOptions() options.backend = "Interpreter" spec = GlowCompileSpec() spec.set(inputs, options) lowered_module = torch_glow.to_glow(traced_m, {"forward": spec}) return lowered_module
def run_model(m, input, randomize): torch_glow.disableFusionPass() traced_m = torch.jit.trace(m, input) input_meta = InputMeta() input_meta.set_same_as(input) inputs = [input_meta] options = CompilationOptions() options.backend = "Interpreter" options.randomize_constants = randomize spec = GlowCompileSpec() spec.set(inputs, options) glow_m = torch_glow.to_glow(traced_m, {"forward": spec}) return glow_m.forward(input)
def run_model(m, input, randomize): if randomize: torch_glow.enable_randomize_constants() else: torch_glow.disable_randomize_constants() torch_glow.disableFusionPass() traced_m = torch.jit.trace(m, input) spec = torch.classes.glow.GlowCompileSpec() spec.setBackend("Interpreter") sim = torch.classes.glow.SpecInputMeta() sim.setSameAs(input) spec.addInputs([sim]) glow_m = torch_glow.to_glow(traced_m, {"forward": spec}) return glow_m.forward(input)
def run_to_glow(m, x): """Trace the model m with input x and call to_glow""" traced_m = torch.jit.trace(m, (x)) spec = torch_glow.CompilationSpec() spec.get_settings().set_glow_backend("Interpreter") compilation_group = torch_glow.CompilationGroup() spec.compilation_groups_append(compilation_group) input_spec = torch_glow.InputSpec() input_spec.set_same_as(x) compilation_group.input_sets_append([input_spec]) lowered_module = torch_glow.to_glow(traced_m, spec) return lowered_module
def test_to_glow_multiple_groups_and_input_sets(self): x1 = torch.randn(1, 4) y1 = torch.randn(2, 4) x2 = torch.randn(1, 2) y2 = torch.randn(5, 2) x3 = torch.randn(7) y3 = torch.randn(3, 7) mod = Foo() scripted_mod = torch.jit.script(mod) x1_y1_set = torch_glow.input_specs_from_tensors([x1, y1]) x2_y2_set = torch_glow.input_specs_from_tensors([x2, y2]) x3_y3_set = torch_glow.input_specs_from_tensors([x3, y3]) # Create two CompilationGroup, first one contains two input sets # and the second CompilationGroup has the third input set spec = torch_glow.CompilationSpec() spec.get_settings().set_glow_backend("Interpreter") compilation_group_1 = torch_glow.CompilationGroup() compilation_group_2 = torch_glow.CompilationGroup() spec.compilation_groups_append(compilation_group_1) spec.compilation_groups_append(compilation_group_2) compilation_group_1.input_sets_append(x1_y1_set) compilation_group_1.input_sets_append(x2_y2_set) compilation_group_2.input_sets_append(x3_y3_set) lowered_module = torch_glow.to_glow(scripted_mod, spec) torch_res1 = mod(x1, y1) torch_res2 = mod(x2, y2) torch_res3 = mod(x3, y3) glow_res1 = lowered_module(x1, y1) glow_res2 = lowered_module(x2, y2) glow_res3 = lowered_module(x3, y3) assert torch.allclose(torch_res1, glow_res1) assert torch.allclose(torch_res2, glow_res2) assert torch.allclose(torch_res3, glow_res3)
def test_serialization(self): with torch.no_grad(): x = torch.randn([1, 4, 4, 4], dtype=torch.float32) y = torch.randn([1, 4, 4, 4], dtype=torch.float32) model = Bar() model = torch.jit.trace(model, (x, y)) spec = torch_glow.CompilationSpec() spec_settings = spec.get_settings() spec_settings.set_glow_backend("NNPI") # Enabled the serialize in this spec spec_settings.set_enable_serialize(True) compilation_group = torch_glow.CompilationGroup() compilation_group_settings = compilation_group.get_settings() compilation_group_settings.set_replication_count(1) compilation_group_settings.backend_specific_opts_insert( "NNPI_IceCores", "1") compilation_group.input_sets_append( torch_glow.input_specs_from_tensors([x, y])) spec.compilation_groups_append(compilation_group) torch_glow.disableFusionPass() torch_glow.enable_convert_to_fp16() # Enable global serialize # then compile(serialize) the model and save it torch_glow.enable_dump_serialized_model() glow_mod = torch_glow.to_glow(model, spec) res1 = glow_mod(x, y) torch.jit.save(glow_mod, "/tmp/serialize_to_glow.pt") # Enable global deserialize and disable serialize # and load(deserialize) the model to loaded_glow_mod torch_glow.enable_deserialize() torch_glow.disable_dump_serialized_model() loaded_glow_mod = torch.jit.load("/tmp/serialize_to_glow.pt") res2 = loaded_glow_mod(x, y) assert torch.allclose(res1, res2, 1e-5, 1e-5)
def test_input_spec(self): """Test setting quantized and non-quantized input specs.""" with torch.no_grad(): a = torch.tensor([[0.1]]) b = torch.tensor([[0.1]]) mod = TestModule() traced_model = torch.jit.trace(mod, (a, b)) ref_result = traced_model(a, b) # test non-quantized input glow_mod = torch_glow.to_glow(traced_model, get_compilation_spec((a, b))) glow_result = glow_mod(a, b) self.assertTrue(torch.allclose(ref_result, glow_result)) # test quantized input add_inputs = torch_glow.get_submod_inputs(mod, "add", (a, b)) glow_mod = torch_glow.to_glow_selective( traced_model, {"add": get_compilation_spec(add_inputs)} ) glow_result = glow_mod(a, b) self.assertTrue(torch.allclose(ref_result, glow_result))
def run_model(m, input, randomize): torch_glow.disableFusionPass() traced_m = torch.jit.trace(m, input) if randomize: torch_glow.enable_randomize_constants() else: torch_glow.disable_randomize_constants() spec = torch_glow.CompilationSpec() spec.get_settings().set_glow_backend("Interpreter") compilation_group = torch_glow.CompilationGroup() spec.compilation_groups_append(compilation_group) input_spec = torch_glow.InputSpec() input_spec.set_same_as(input) compilation_group.input_sets_append([input_spec]) glow_m = torch_glow.to_glow(traced_m, {"forward": spec}) return glow_m(input)
def test_to_glow_tuple_output(self): a = torch.randn(4) model = Foo() torch_resA = model(a) (tx, ty) = torch_resA metaA = torch_glow.InputMeta() metaA.set_same_as(a) inputA = [metaA] options = torch_glow.CompilationOptions() options.backend = "Interpreter" specA = torch_glow.GlowCompileSpec() specA.set(inputA, options) scripted_mod = torch.jit.script(model) lowered_mod = torch_glow.to_glow(scripted_mod, [specA]) glow_resA = lowered_mod(a) (gx, gy) = glow_resA assert torch.allclose(tx, gx) assert torch.allclose(ty, gy)
def lower_and_write_to_onnx_helper(self, ModType, onnx_prefix): x = torch.randn(1, 3, 8, 8) model = create_model(x, ModType) spec = torch_glow.CompilationSpec() spec.get_settings().set_glow_backend("Interpreter") compilation_group = torch_glow.CompilationGroup() spec.compilation_groups_append(compilation_group) input_spec = torch_glow.InputSpec() input_spec.set_same_as(x) compilation_group.input_sets_append([input_spec]) scripted_mod = torch.jit.trace(model, x) torch_glow.enable_write_to_onnx() torch_glow.set_onnx_file_name_prefix(onnx_prefix) torch_glow.enable_write_without_randomize() lowered_model = torch_glow.to_glow(scripted_mod, {"forward": spec}) # Run Glow model g = lowered_model(x) # Run reference model t = model(x) self.assertEqual(type(g), type(t)) self.assertEqual(len(g), len(t)) for (gi, ti) in zip(g, t): self.assertTrue(torch.allclose(gi, ti)) assert os.path.exists(onnx_prefix + ".onnxtxt") onnx_files = glob.glob(onnx_prefix + "*.onnx*") for f in onnx_files: os.remove(f)
def infer_nnpi(model, device, data_type, input_size, output_size, batch_size, args): import torch_glow # Detailed structure for spec can be found at https://fburl.com/diffusion/79q4efud # Create compilation spec spec = torch_glow.CompilationSpec() spec.get_settings().set_glow_backend("NNPI") # Create compilation group and update settings. # Compilation group contains compilation specific information like # fp16 settings, enableRemoveMutation, anything that changes # the Glow graph compiled for example. compilation_group = torch_glow.CompilationGroup() compilation_group_settings = compilation_group.get_settings() compilation_group_settings.set_convert_to_fp16(True) compilation_group_settings.set_replication_count(1) compilation_group_settings.backend_specific_opts_insert( "NNPI_IceCores", "1") data = torch.randn(batch_size, input_size) # Create input spec and add it into compilation group. # This is used for shape inference when lowering the model to Glow. data_spec = torch_glow.InputSpec() data_spec.set_same_as(data) compilation_group.input_sets_append([data_spec]) spec.compilation_groups_append(compilation_group) traced_model = torch.jit.trace(model, (data)) lowered_model = torch_glow.to_glow(traced_model, spec) start_time = time.time() for i in range(args.steps + args.warmups): lowered_model(data) if i < args.warmups: start_time = time.time() return time.time() - start_time