npu_model

Applicability

Product	Supported
Atlas A3 training products/Atlas A3 inference products	√
Atlas A2 training products/Atlas A2 inference products	√
Atlas 200I/500 A2 inference products	x
Atlas inference products	x
Atlas training products	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 the PyTorch scenario, if all computations are performed on the device and the input and output are also on the device, data needs to be copied from the device to the host. After the execution, PyTorch moves the data to the device, affecting the execution performance. npu_model can be used to optimize the execution so that moving data is not required (that is, the data is directly offloaded to the device for execution).

Prototype

1	Decorator@npu_model

Parameters

Parameter

Data Type

Value Description

optimize_level

int

1: In the PyTorch scenario, the UDF NN engine is used to offload the input and output data to the device for execution. The default value is 1.
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 parameter must be used together with input_descs.

Note that this configuration 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

This hook function is used to set up the Python environment or perform 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.

The DfException exception is reported in abnormal scenarios. You can obtain error_code and message in DfException to view the error code and error information. 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