Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
multi_margin_loss.py 2.38 KB
Newer Older
yanjl1's avatar
Initial  
yanjl1 committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64 
import hipdnn
import torch


def build_multi_margin_loss_graph(
    hipdnn_handle, torch_tensor_input, torch_tensor_target, torch_tensor_weight, hipdnn_data_type
):
    graph = hipdnn.pygraph(
        handle=hipdnn_handle,
        io_data_type=hipdnn_data_type,
        intermediate_data_type=hipdnn.data_type.FLOAT,
        compute_data_type=hipdnn.data_type.FLOAT,
        name="multi_margin_loss",
    )
    hipdnn_tensor_input = graph.tensor_like(torch_tensor_input)
    hipdnn_tensor_target = graph.tensor_like(torch_tensor_target)
    hipdnn_tensor_weight = graph.tensor_like(torch_tensor_weight)
    hipdnn_tensor_output = graph.multi_margin_loss(
        input=hipdnn_tensor_input,
        target=hipdnn_tensor_target,
        weight=hipdnn_tensor_weight,
        p=1,
        margin=1.0,
        reduction=hipdnn.reduction_mode.AVG,
        name="multi_margin_loss",
    )
    hipdnn_tensor_output.set_output(True)
    graph.build(hipdnn_handle)
    return (
        graph,
        hipdnn_tensor_input,
        hipdnn_tensor_target,
        hipdnn_tensor_weight,
        hipdnn_tensor_output,
    )


if __name__ == "__main__":
    batch, num_classes = 4, 10
    hipdnn_data_type = hipdnn.data_type.FLOAT
    torch_data_type = torch.float32
    torch_tensor_input = torch.rand(batch, num_classes, dtype=torch_data_type, device="cuda")
    torch_tensor_target = torch.randint(0, num_classes, (batch,), dtype=torch.int64, device="cuda")
    torch_tensor_weight = torch.ones(num_classes, device="cuda")
    hipdnn_handle = hipdnn.create_handle()
    graph, hipdnn_tensor_input, hipdnn_tensor_target, hipdnn_tensor_weight, hipdnn_tensor_output = (
        build_multi_margin_loss_graph(
            hipdnn_handle,
            torch_tensor_input,
            torch_tensor_target,
            torch_tensor_weight,
            hipdnn_data_type,
        )
    )
    torch_tensor_output = torch.empty(batch, dtype=torch_data_type, device="cuda")
    variant_pack = {
        hipdnn_tensor_input: torch_tensor_input.data_ptr(),
        hipdnn_tensor_target: torch_tensor_target.data_ptr(),
        hipdnn_tensor_weight: torch_tensor_weight.data_ptr(),
        hipdnn_tensor_output: torch_tensor_output.data_ptr(),
    }
    workspace = torch.empty(graph.get_workspace_size(), dtype=torch.uint8, device="cuda")
    graph.exec(variant_pack=variant_pack, workspace=workspace.data_ptr())
    print("multi_margin_loss graph execution complete.")