文档首页/ 魔坊（ModelArts）模型训推平台/ 模型训练/ 创建训练作业/ 训练示例代码/ 示例：将本地PyTorch训练代码迁移到ModelArts
更新时间：2026-08-13 GMT+08:00
查看PDF
分享

示例：将本地PyTorch训练代码迁移到ModelArts

方案概述

本实践旨在指导开发者将本地已调试完成的PyTorch训练脚本，改造为可在ModelArts训练作业中运行的标准形态，内容涵盖数据路径改造、运行环境准备、单机多卡/多机多卡分布式改造，以及从GPU迁移至昇腾NPU的代码适配四个部分，并分别给出改造前后的代码示例。

代码必须先在本地调试好，再上传华为云平台进行训练，否则会出现各种难以解释的问题，因此本实践建议严格按照“本地小规模验证→训练作业联调→正式训练作业”的顺序推进迁移。

数据与输出路径改造

  1. 改造背景

    本地训练脚本通常直接读写本地磁盘路径，而ModelArts训练作业运行在容器中，数据与模型产出需要经由OBS对象存储中转。因此改造的第一步是将脚本中的本地路径替换为可由平台注入的data_url（训练输入）与train_url（训练输出）参数。在创建训练作业时，需要单击“增加训练输入”，设置训练输入的“参数名称”为“data_url”，并设置数据存储位置为OBS目录；对于输出，则单击“增加训练输出”，设置训练输出的“参数名称”为“train_url”，同样对应一个OBS目录。

  2. 代码示例。
    • 迁移前代码（设置数据和模型为本地路径）。 
      import argparse
import torch
from torchvision import datasets, transforms
parser = argparse.ArgumentParser()
parser.add_argument('--batch_size', type=int, default=64)
args = parser.parse_args()
# 本地的数据路径
train_data = datasets.MNIST('/home/user/data/mnist', train=True,
                             download=False, transform=transforms.ToTensor())
train_loader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size)
# ... 训练过程 ...
# 本地的模型保存路径
torch.save(model.state_dict(), '/home/user/output/model.pth')
    • 迁移后（参数化data_url / train_url，OBS与容器本地路径自动映射）。 
      import argparse
import os
import torch
from torchvision import datasets, transforms
parser = argparse.ArgumentParser()
parser.add_argument('--batch_size', type=int, default=64)
# ModelArts 会将 data_url 对应的 OBS 数据下载到容器本地路径后传入
parser.add_argument('--data_url', type=str, default='/home/work/data')
# ModelArts 会自动将 train_url 对应容器本地目录中的内容回传至 OBS
parser.add_argument('--train_url', type=str, default='/home/work/output')
args = parser.parse_args()
train_data = datasets.MNIST(args.data_url, train=True,
                             download=False, transform=transforms.ToTensor())
train_loader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size)
# ... 训练过程 ...
os.makedirs(args.train_url, exist_ok=True)
torch.save(model.state_dict(), os.path.join(args.train_url, 'model.pth'))

      如果用户在自己的训练脚本中要创建新的目录或文件，请在inputs或者outputs中指定的local_path等目录中创建，避免因权限或路径隔离导致写入失败。

      如果使用ModelArts SDK方式在Notebook中调试后再提交远程训练，输出配置遵循相同约定：local_path为Notebook中的本地路径，训练脚本需要将输出的模型或其他数据保存在该目录下；obs_path为OBS目录，SDK会将local_path中的模型文件自动上传到这里。本文示例代码是以PyTorch为例编写的，不同的AI框架之间整体流程完全相同，仅需修改framework_type参数值即可，切换到其他框架时无需重写迁移逻辑。

运行环境（镜像）改造

  1. 选型原则。

    如果本地PyTorch/CUDA版本与ModelArts预置镜像匹配，建议直接复用预置镜像，例如引擎及版本下拉框中选择PyTorch，pytorch_1.8.0-cuda_10.2-py_3.7-ubuntu_18.04-x86_64。因为这些镜像经过充分的功能验证，并且已经预置了很多常用的安装包，只有在预置镜像无法满足特殊依赖时才需要制作自定义镜像。

  2. 自定义镜像构建规范

    当依赖复杂、需要制作自定义镜像时，应遵循以下规范：容器镜像的大小建议小于15G；建议通过开源的官方镜像来构建，例如PyTorch的官方镜像；建议容器分层构建，单层容量不要超过1G、文件数不大于10万个，分层时先构建不常变化的层，例如先OS，再cuda驱动，再Python，再pytorch，再其他依赖包；如果训练数据和代码经常变动，则不建议把数据、代码放到容器镜像里，避免频繁地构建容器镜像。

    如果需将本地conda环境迁移至容器（容器已能满足环境隔离需求，不建议在容器内再创建多个conda env），可采用打包迁移方式：

    # 在本地/云主机上，基于想要迁移的base环境创建一个名为pytorch的conda环境
conda create --name pytorch --clone base
pip install conda-pack
# 将pytorch env打包生成pytorch.tar.gz
conda pack -n pytorch -o pytorch.tar.gz
  3. 自定义镜像下的训练作业配置

    使用自定义镜像时，训练作业的镜像地址、代码目录、日志路径分别指向SWR与OBS，例如镜像地址填写为“swr.cn-north-4.myhuaweicloud.com/deep-learning/pytorch:1.8.1-cuda11.1”，代码目录设置为OBS中存放启动脚本文件的目录，训练代码会被自动下载至训练容器的${MA_JOB_DIR}/demo-code目录中。

单机多卡 / 多机多卡分布式改造

  1. 改造思路。

    本地若仅使用单卡训练，迁移到ModelArts做分布式扩展时，推荐采用DistributedDataParallel（DDP）而非DataParallel（DP），因为DDP能够启动多进程进行运算，从而大幅度提升计算资源的利用率，可以基于torch.distributed实现真正的分布式计算。DP与DDP的代码差异示例：

    import torch
class Net(torch.nn.Module):
    pass
model = Net().cuda()
### DataParallel Begin ###
model = torch.nn.DataParallel(Net().cuda())
### DataParallel End ###
  2. DDP改造代码示例。

    以下以官方对resnet18在cifar10数据集上的分类任务改造为例，训练代码中包含三部分入参，分别为训练基础参数、分布式参数和数据相关参数，其中分布式参数由平台自动入参，无需自行定义：

    • 改造前（单卡训练，核心训练入口） 
      import torch
from torch import nn, optim
def main():
    model = ResNet18()
    model.cuda()
    optimizer = optim.SGD(model.parameters(), lr=0.01)
    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
    for epoch in range(epochs):
        for data, label in train_loader:
            data, label = data.cuda(), label.cuda()
            optimizer.zero_grad()
            loss = nn.CrossEntropyLoss()(model(data), label)
            loss.backward()
            optimizer.step()
    • 改造后（新增分布式初始化、Sampler、DDP包装，标注为分布式改造点） 
      import argparse
import torch
import torch.distributed as dist
from torch import nn, optim
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
parser = argparse.ArgumentParser()
parser.add_argument('--init_method', default=None, help='tcp_port')
parser.add_argument('--rank', type=int, default=0, help='index of current task')
parser.add_argument('--world_size', type=int, default=1, help='total number of tasks')
args, unknown = parser.parse_known_args()
def main():
    ### 分布式改造，初始化进程组 ###
    dist.init_process_group(backend='nccl', init_method=args.init_method,
                             rank=args.rank, world_size=args.world_size)
    torch.cuda.set_device(args.rank % torch.cuda.device_count())
    ### 分布式改造结束 ###
    model = ResNet18().cuda()
    ### 分布式改造，使用DDP包装模型 ###
    model = torch.nn.parallel.DistributedDataParallel(model)
    ### 分布式改造结束 ###
    optimizer = optim.SGD(model.parameters(), lr=0.01)
    ### 分布式改造，使用DistributedSampler对数据切分 ###
    train_sampler = DistributedSampler(train_dataset)
    train_loader = DataLoader(train_dataset, batch_size=64, sampler=train_sampler)
    ### 分布式改造结束 ###
    for epoch in range(epochs):
        train_sampler.set_epoch(epoch)   # 分布式改造，确保每轮shuffle不同
        for data, label in train_loader:
            data, label = data.cuda(), label.cuda()
            optimizer.zero_grad()
            loss = nn.CrossEntropyLoss()(model(data), label)
            loss.backward()
            optimizer.step()

      以上代码支持多节点分布式训练，同时兼容CPU和GPU分布式训练环境，用户可以通过注释掉代码中的分布式改造点，轻松切换为单节点单卡训练模式。其中init_method参数值会包含主节点的ip和端口，由平台自动入参，不需要用户输入主节点的ip和端口；当资源规格为单机多卡时，需要在创建训练作业时指定超参world_size和rank，若资源规格为多机时（训练作业计算节点个数大于1）则无需设置，world_size和rank超参由平台自动注入。

  3. 大数据集的存储优化。

    对于数据量较大的分布式训练场景，建议先通过obsutil工具将数据集传到OBS桶后，再将数据集迁移至SFS（弹性文件服务）以提升IO性能，示例命令为：

    # 将OBS的代码传到SFS中 
./obsutil cp obs://your_bucket/YOLOX/ /mnt/sfs_turbo/code/ -f -r

GPU训练迁移至昇腾NPU训练

如果ModelArts侧使用的是昇腾NPU算力，除上述云化改造外，还需要完成硬件适配层面的代码迁移。需要注意的是，NPU（Neural Network Processing Unit）和GPU在构造结构上存在差异，因此迁移过程并不是完全平替的关系，虽然在表达层可以通过torch.cuda和torch.npu的形式来替代，但真实的算子下发、显存管理、集合通信等仍存在差异。

  1. 安装Ascend Extension for PyTorch（torch_npu）

    PyTorch官方并不直接支持昇腾的后端，仅直接支持CUDA和AMD ROCm，因此原生PyTorch的GPU训练代码无法直接在昇腾设备运行；PyTorch 2.1版本提供了新硬件适配的插件机制，通过安装Ascend Extension for PyTorch插件后，即可直接使用PyTorch的表达层运行在NPU设备上。安装校验：

    python3 -c "import torch;import torch_npu;print(torch_npu.npu.is_available())"
  2. 自动迁移（推荐，简单场景优先尝试）

    如果没有用到GPU的高阶能力，例如自定义算子、直接操作GPU显存等操作，可以直接使用自动迁移，仅需在训练入口脚本import torch之后添加两行代码：

    • 迁移前： 
      import torch
import torch.nn as nn
# ... 正常的GPU训练代码 ...
    • 迁移后： 
      import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu   # 自动映射cuda API到npu
import torch.nn as nn
# ... 训练代码无需修改，torch.cuda相关调用会被自动转换为torch.npu对应操作 ...
  3. 手动迁移（自动迁移失败或含GPU高阶能力时）

    如果自动迁移后训练仍报错，则需要手动将CUDA相关接口逐一替换为NPU接口。核心改造点包括设备指定、CUDA接口替换、以及分布式通信后端切换：

    • 设备指定方式迁移前后对比。 
      # 迁移前
device = torch.device('cuda:{}'.format(args.gpu))
torch.cuda.set_device(args.gpu)
# 迁移后
device = torch.device('npu:{}'.format(args.gpu))
torch_npu.npu.set_device(args.gpu)
    • 常用CUDA接口替换 
      # 迁移前
torch.cuda.is_available()
model.cuda(args.gpu)
images = images.cuda(args.gpu, non_blocking=True)
target = target.cuda(args.gpu, non_blocking=True)
# 迁移后
torch_npu.npu.is_available()
model.npu(args.gpu)
images = images.npu(args.gpu, non_blocking=True)
target = target.npu(args.gpu, non_blocking=True)
    • 分布式通信后端切换（多卡场景，在单卡迁移的3个修改要点之外还需切换通信后端）： 
      # 迁移前（GPU使用nccl）
dist.init_process_group(backend='nccl', init_method="tcp://127.0.0.1:port",
                         ......, rank=args.rank)
# 迁移后（NPU使用hccl）
dist.init_process_group(backend='hccl', init_method="tcp://127.0.0.1:port",
                         ......, rank=args.rank)

    昇腾NPU平台不支持torch.nn.DataParallel接口，若训练脚本中使用了该接口，需要手动修改为torch.nn.parallel.DistributedDataParallel接口执行多卡训练。

  4. 精度与性能核验。

    迁移完成后建议使用msprobe工具对比标杆（GPU/CPU）环境和昇腾环境上运行训练时的差异点，主要包括精度预检、精度比对和梯度监控等功能，排查精度问题时该工具会自动将torch.nn.functional.dropout、torch.nn.Dropout等接口参数p（丢弃概率）置为0以消除随机性干扰。性能方面，PyTorch在昇腾AI处理器上是以算子为粒度进行调用（OP-based），性能优化的总体原则为减少Host算子下发时间、减少Device算子执行时间；ModelArts提供的MA-Advisor性能自动诊断工具可自动扫描profiling数据并给出调优建议，实测大约可提升10%~30%的性能。

    若自动迁移或手动改造中遇到无法解决的报错，可以首先在昇腾社区论坛以及Gitee的PyTorch Issues中查找线索，如果还无法解决，可以通过提交工单的形式从华为云ModelArts入口进行咨询以及求助对应的专业服务。

迁移检查清单

  • 训练脚本已将本地路径改造为data_url/train_url参数化输入输出。
  • 已确认镜像版本（预置或自定义）与本地PyTorch/CUDA版本兼容。
  • 大数据集场景已规划OBS→SFS的数据中转路径。
  • 分布式场景已完成DDP改造，并区分单机多卡/多机多卡的超参注入方式。
  • 如目标为NPU，已完成torch_npu安装验证并优先尝试transfer_to_npu自动迁移。
  • 自动迁移失败场景已针对设备接口、通信后端等逐项手动替换。
  • 已用msprobe等工具完成精度比对，用MA-Advisor完成性能诊断。
父主题： 训练示例代码

相关文档