MindTorch基于MindSpore的实现接口功能映射和执行,因此天然可继承MindSpore框架的静态图、自动并行等优势。本章节描述已完成网络基本功能迁移后使用静态图等功能进行加速训练。

静态图加速

MindSpore框架的执行模式有两种:动态图(PyNative)模式和静态图(Graph)模式:

  • 动态图模式下,程序按照代码的编写顺序执行,在执行正向过程中根据反向传播的原理,动态生成反向执行图。动态图模式方便编写和调试神经网络模型。

  • 静态图模式下,程序在编译执行时先生成神经网络的图结构,然后再执行图中涉及的计算操作。静态图模式利用图优化等技术对执行图进行更大程度的优化,因此能获得较好的性能,但是执行图是从源码转换而来,因此在静态图下不是所有的Python语法都能支持。 更多详细信息请参考MindSpore动静统一机制介绍

目前MindTorch默认支持PyNative模式,请首先在PyNative模式下完成功能调试。如果想调用静态图模式进行训练加速,再尝试切换到Graph模式执行。下面介绍两种切换静态图的方式:

方式一:采用即时编译装饰器jit,使能部分函数粒度表达模块以静态图模式执行。

import mindspore as ms
ms.set_context(jit_syntax_level=ms.STRICT)
@ms.jit
def mul(x, y):
    return x * y

方式二:全局设置Graph模式,更适合基于Module表达。

import mindspore as ms
ms.set_context(mode=ms.GRAPH_MODE)
ms.set_context(jit_syntax_level=ms.STRICT)

由于Graph模式下不是所有的Python语法都能支持,通过上面两种方式切换到Graph模式后部分网络可能会出现语法不支持情况,推荐调整日志级别export MSA_LOG=2协助调试分析,根据报错信息对代码进行相应调整,当前主要体现在in-place类型操作和部分语法用法限制,具体可参考静态图语法支持

原生框架并行训练

分布式并行训练可以降低对内存、计算性能等硬件的需求,是进行训练的重要优化手段。目前MindTorch中对标torch.distributed相关分布式接口还在开发中,如果用户想要使用分布式训练进行加速训练,可以将torch.distributed相关接口替换成MindSpore提供的更简单易用的高阶API。MindTorch基于MindSpore分布式并行能力提供两种并行模式:

  • 数据并行:对数据进行切分的并行模式,一般按照batch维度切分,将数据分配到各个计算单元中,进行模型计算。

  • 自动并行:融合了数据并行、算子级模型并行的分布式并行模式,可以自动建立代价模型,找到训练时间较短的并行策略,为用户选择合适的并行模式。仅支持在图模式下使用。

相关机制请参考MindSpore原生分布式并行架构

1.1 数据并行

import mindtorch.torch as torch
from mindtorch.torch.utils.data import DataLoader, DistributedSampler
from mindspore.communication import init
import mindspore as ms

...
init("hccl")   # 初始化通信环境:"hccl"---Ascend, "nccl"---GPU, "mccl"---CPU
ms.set_auto_parallel_context(parallel_mode=ms.ParallelMode.DATA_PARALLEL) # 配置数据并行模式
torch.manual_seed(1)     # 设置随机种子,使得每张卡上权重初始化值一样,便于收敛

train_images = datasets.CIFAR10('./', train=True, download=True, transform=transform)
sampler = DistributedSampler(train_images)     # 分布式数据处理
train_data = DataLoader(train_images, batch_size=32, num_workers=2, drop_last=True, sampler=sampler)

...

def forward_fn(data, label):
    logits = net(data)
    loss = criterion(logits, label)
    return loss, logits

grad_fn = ms.ops.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=True)
grad_reducer = nn.DistributedGradReducer(optimizer.parameters)  # 定义分布式优化器

def train_step(data, label):
    (loss, _), grads = grad_fn(data, label)
    grads = grad_reducer(grads)   # 梯度聚合
    optimizer(grads)
    return loss

net.train()
for i in range(epochs):
    for inputs, target in train_data:
        res = train_step(inputs, target)
...

1.2 自动并行

import mindtorch.torch as torch
from mindtorch.torch.utils.data import DataLoader, DistributedSampler
from mindspore.communication import init
import mindspore as ms

...
ms.set_context(mode=ms.GRAPH_MODE, jit_syntax_level=True)   # 自动并行仅支持静态图模式
init("hccl")   # 初始化通信环境:"hccl"---Ascend, "nccl"---GPU, "mccl"---CPU
ms.set_auto_parallel_context(parallel_mode=ms.ParallelMode.AUTO_PARALLEL, search_mode="recursive_programming")  # 配置自动并行模式
torch.manual_seed(1)   # 设置随机种子,使得每张卡上权重初始化值一样,便于收敛

train_images = datasets.CIFAR10('./', train=True, download=True, transform=transform)
sampler = DistributedSampler(train_images)     # 分布式数据处理
train_data = DataLoader(train_images, batch_size=32, num_workers=2, drop_last=True, sampler=sampler)

...

def forward_fn(data, label):
    logits = net(data)
    loss = criterion(logits, label)
    return loss, logits

grad_fn = ms.ops.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=True)

@ms.jit
def train_step(data, label):
    (loss, _), grads = grad_fn(data, label)
    optimizer(grads)
    return loss

net.train()
for i in range(epochs):
    for inputs, target in train_data:
        res = train_step(inputs, target)
...

MindSpore自动并行功能目前在实验性阶段,仅支持部分场景。如果在使用过程中出现不支持的报错信息,可以通过ISSUE反馈。