7.1.6. Example: Explicitly Transferring Data Between Host And Device
Example: Inference With A Model-based sample program demonstrating explicit host-to-device data transfer APIs
Execution Method
$ cd /opt/pfn/pfcomp/codegen/examples/
$ ./exec_with_env.sh python3 explicit_data_transfer_api.py
Expected Output
Inference results using a randomly initialized model (should be identical to those from
infer.py)
tensor([[-0.3188, 0.6279, -0.5000, -1.2148],
[-0.3188, 0.6279, -0.5000, -1.2148],
[-0.3188, 0.6279, -0.5000, -1.2148],
[-0.3188, 0.6279, -0.5000, -1.2148]])
Related Links
Sample Program
1import torch
2from mlsdk import Context, MNDevice, set_tensor_name_in_module, storage
3
4torch.manual_seed(0)
5
6
7def run_explicit_data_transfer():
8 device = MNDevice("mncore2:auto")
9 context = Context(device)
10 Context.switch_context(context)
11
12 model = torch.nn.Linear(4, 4)
13 model.eval()
14 set_tensor_name_in_module(model, "model")
15 for p in model.parameters():
16 context.register_param(p)
17 for b in model.buffers():
18 context.register_buffer(b)
19
20 def infer(input: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
21 x = input["x"]
22 y = model(x)
23 return {"out": y}
24
25 sample = {"x": torch.randn(4, 4)}
26
27 compiled_infer = context.compile(
28 infer,
29 sample,
30 storage.path("/tmp/proxy_transfer"),
31 )
32
33 input_proxies_allocated = compiled_infer.allocate_input_proxy()
34 input_proxies_allocated["x"].load_from(torch.ones(4, 4), clone=False)
35
36 for model_param in model.parameters():
37 model_param_proxy = context.get_registered_value_proxy(model_param)
38 model_param_proxy.load_from(model_param)
39
40 result = compiled_infer(input_proxies_allocated)
41 result_on_cpu = result["out"].cpu()
42 print(result_on_cpu)
43
44
45if __name__ == "__main__":
46 run_explicit_data_transfer()