npu_model

Applicability

Product	Supported
Atlas A3 training product/Atlas A3 inference product	√
Atlas A2 training product/Atlas A2 inference product	√
Atlas 200I/500 A2 inference product	x
Atlas inference product	x
Atlas training product	x

Function Description

If the UDF is deployed on the host, data needs to be copied from the device to the local host for computation. In PyTorch scenarios, when all computations, inputs, and outputs reside on the device, data still has to be copied from the device to the host during execution, followed by PyTorch transferring the data back to the device, which degrades performance. The npu_model API optimizes this process by eliminating data transfer overhead, allowing computations to be directly offloaded to the device for execution.

Prototype

1	Decorator@npu_model

Parameters

Parameter

Data Type

Description

optimize_level

int

1 (default): In PyTorch scenarios, the UDF NN engine is used to offload the input and output data to the device for execution.
2: Compiles the PyTorch model into a graph, exports the graph as an NN model, and optimizes the model to an NPU model for loading and execution. This option must be used together with input_descs.

Note that this configuration item takes effect when a class is decorated. It cannot be configured for a decorated function.

input_descs

[TensorDesc]

Describes the input tensor of the graph exported from Torch when optimize_level is set to 2. The following is an example:

input_descs=[TensorDesc(dtype = df.DT_INT64, shape = [2,1,4]),TensorDesc(dtype =
df.DT_FLOAT, shape = [2,1,4])],

If a dimension in the shape is negative, the input is dynamic. A dynamic graph will be exported using npu_model.

num_returns

int

Indicates the number of outputs when a decorator is used. If this parameter is not set, only one value is returned by default. You can either use this parameter to set the number of returned values or use type annotations to set the number and type of returned values.

resources

dict

Indicates the resource information required by the current function, which can be memory, num_cpus, or num_npus. The unit of memory is MB. num_npus indicates the number of NPUs to be used. num_npus is reserved and can only be set to 1. Example: {"memory": 100, "num_cpus": 1, "num_npus": 1}

env_hook_func

function

Sets up the Python environment or performs the import operation before the Python UDF is initialized.

visible_device_enable

bool

After this function is enabled, the UDF process automatically sets ASCEND_RT_VISIBLE_DEVICES based on the configured num_npus and calls the get_running_device_id API to obtain the corresponding logical ID. Currently, num_npus supports only 1. Therefore, the get_running_device_id API returns 0 in this scenario.

Returns

The decorated functions are returned in normal scenarios.

A DfException is thrown upon exceptions. You can catch DfException and retrieve its error_code and message attributes to check the specific error code and error details. For details, see DataFlow Error Codes.

Examples

@df.npu_model(optimize_level=1)
class FakeModel1(nn.Module):
    def __init__(self):
        super().__init__()

    # Simulate model inference.
    @df.method()
    def forward(self, input_image):
        return F.interpolate(input_image, size=(256, 256), mode='bilinear')

@df.npu_model(optimize_level=1, input_descs=[df.TensorDesc(dtype=df.DT_FLOAT, shape=[1, 3, 768, 768])])
class FakeModel2(nn.Module):
    def __init__(self):
        super().__init__()
        self.mean = 0.5
        self.std = 0.5

    # Simulate model inference.
    @df.method()
    def forward(self, input_image):
        return (input_image - self.mean) / self.std

@df.npu_model()
def preprocess(input_image):
    # Simulate image cropping.
    transform = transforms.Compose([transforms.CenterCrop(512)])
    return transform(input_image)

@df.npu_model()
def postprocess(input_image):
    mean = 0.5
    std = 0.5
    img = input_image * std + mean
    return F.interpolate(img, size=(512, 512), mode='bilinear')

Constraints

The torch_npu package of the corresponding Python version must be installed.
The input and output must be NPU tensors.
One group of inputs corresponds to one group of outputs. Streaming input and output are not supported.

Parent topic: Graph Construction APIs