Operator Usage Guide (ATB Python APIs)

  • For details about the complete executable example, see Graph Example.
  • Before the running, you need to install the torch_atb plugin. The installation method is as follows:
    ./Ascend-cann-nnal_${version_info}_linux-aarch64.run --install --torch_atb
  1. Import the ATB Python API module.

    Import the ATB Python API module by using the following code:

    import torch
    import torch_atb
  2. Create an operator object instance.
    • Single-operator
      1. Create a parameter object.

        Instantiate the parameter structure based on the operator to be created. For details about the definition of the parameter API, see OpParam. The following uses linear as an example.

        linear_param = torch_atb.LinearParam()
        linear_param.has_bias = False
      2. Create an operator object.
        op = torch_atb.Operation(linear_param)
    • Graph operator

      Create a graph operator based on the configured mappings between tensors of composite operators by using tensorName and the designed graph operator structure.

      Use the graph constructor to construct a graph.

      • Graph construction method 1:
         1
         2
         3
         4
         5
         6
         7
         8
         9
        10
        11
        12
        13
        14
            builder = torch_atb.Builder("Graph")
            x = builder.add_input("x")
            y = builder.add_input("y")
            elewise_add = torch_atb.ElewiseParam()
            elewise_add.elewise_type = torch_atb.ElewiseParam.ElewiseType.ELEWISE_ADD
            layer1 = builder.add_node([x, y], elewise_add)
            add_out = layer1.get_output(0)
            z = builder.add_input("z")
            builder.reshape(add_out, lambda shape: [1, shape[0] * shape[1]], "add_out_")
            elewise_mul = torch_atb.ElewiseParam()
            elewise_mul.elewise_type = torch_atb.ElewiseParam.ElewiseType.ELEWISE_MUL
            layer2 = builder.add_node(["add_out_", z], elewise_mul)
            builder.mark_output(layer2.get_output(0))
            Graph = builder.build()
        
      • Graph construction method 2:
        graph = torch_atb.GraphBuilder("Graph") \ # Set the name of the graph operator.
                 .set_input_output(["x", "y", "z"], ["out"]) \ # Specify the input and output tensor names of the graph.
                 .add_operation(elewise_add, ["x", "y"], ["add_out"]) \ # Add the single-operator required in the graph and the used tensor.
                 .reshape("add_out", lambda shape: [1, shape[0] * shape[1]], "add_out_") \ # Use the lambda expression to dynamically update the shape of the output tensor if necessary.
                 .add_operation(elewise_mul, ["add_out_", "z"], ["out"]) \
                 .build() # End graph construction and create a graph operator.
  3. Prepare input data.

    Prepare the input tensor and ensure that the device (NPU or CPU) and data type meet the expectation.

    x = torch.randn(2, 3, dtype=torch.float16).npu()  
    y = torch.randn(2, 3, dtype=torch.float16).npu()
  4. Perform the operation.

    Use the forward method to complete the operation and obtain the output.

    outputs = op.forward([x, y]) # If the graph operator is used, the output is graph.forward([x, y, z]).
    torch.npu.synchronize()