This quick start uses the Atlas 800T A2 training server as an example.
- Install the matching versions of NPU driver firmware and CANN software (Toolkit, ops, and NNAL). For details, see :
- Install the PyTorch framework and the torch_npu plugin. For details, see .
This section provides a simple model migration example using the simplest automatic migration method, helping users quickly experience the process of migrating GPU model scripts to Ascend NPUs. This example is based on a CNN model script for recognizing handwritten digits. The original GPU training code is modified so that it can be migrated to Ascend NPUs for training.
Create a new script train.py and write the following original GPU script code into it.
[object Object]Add the following code to train.py.
- If you are using Atlas training products, you need to enable mixed precision after migration is complete and before training begins, due to their architectural characteristics.
- If you are using Atlas A2 training or Atlas A3 training products, you can choose whether to enable mixed precision.
[object Object]
[object Object]If automatic migration is not enabled, refer to the "" section in the PyTorch Training Model Migration and Tuning Guide.
Enable AMP mixed precision computation. If you are using Atlas A2 training or Atlas A3 training products, you may skip this step.
After defining the model and optimizer, define the GradScaler for AMP functionality.
[object Object]Delete the following original GPU script code.
[object Object]Add the following code to enable AMP.
[object Object]Run the command to start the training script (the script name can be modified as needed).
[object Object]After training completes, a weight file as shown below is generated, indicating that the migration training is successful.