单机单卡场景

环境准备

请参考《Ascend Extension for PyTorch 配置与安装》手册，准备昇腾环境，包括硬件与OS，NPU驱动固件、CANN软件、PyTorch的安装。
根据实际需求准备模型环境，如conda、docker以及三方库依赖。

模型脚本配置

以常见的PyTorch模型，GPU训练为例：

通常模型源码支持GPU单机单卡训练，用户根据模型脚本迁移章节的指导，将模型源码映射至NPU，即可完成模型脚本配置。如果模型源码仅使用了CPU编程，不支持GPU单机单卡训练，用户可使用如下命令，将数据、模型放置NPU上训练：

import torch
import torch_npu

device = torch.device("npu")
# 将数据to npu
for data in dataloder:
  if torch.npu.is_available():
       data.to(device)

# 将模型to npu
model = Model()
if torch.npu.is_available():
     model.to(device)

# 将损失函数to npu
loss_fn = nn.CrossEntropyLoss()
if torch.npu.is_available():
     loss_fn=loss_fn.to(device)

启动脚本配置

执行如下命令启动脚本拉起单机训练（以下参数为举例，用户可根据实际情况自行改动）。

python3 main.py   --batch-size 128 \                          # 训练批次大小 
                  --data_path /home/data/resnet50/imagenet \  # 数据集路径
                  --lr 0.1 \                                  # 学习率 
                  --epochs 90 \                               # 训练迭代轮数 
                  --arch resnet50 \                           # 模型架构 
                  --world-size 1 \ 
                  --rank 0 \          
                  --workers 40 \                              # 加载数据进程数 
                  --momentum 0.9 \                            # 动量   
                  --weight-decay 1e-4 \                       # 权重衰减 
                  --gpu 0                                     # device号, 这里参数名称仍为gpu, 但迁移完成后实际训练设备已在代码中定义为npu

父主题： 模型脚本与启动脚本配置