7.1.6. Example: Explicitly Transferring Data Between Host And Device

Example: Inference With A Model を元に、明示的にホストーデバイス間のデータ転送を行う API を扱うサンプルプログラム

実行方法

$ cd /opt/pfn/pfcomp/codegen/examples/
$ ./exec_with_env.sh python3 explicit_data_transfer_api.py

想定出力

  • ランダムに初期化されたモデルによる推論結果 (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]])

関連リンク

サンプルプログラム

リスト 7.6 /opt/pfn/pfcomp/codegen/MLSDK/examples/explicit_data_transfer_api.py
 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()