Torch-MLU 入门教程

本教程介绍如何使用 Torch-MLU 在寒武纪 MLU 设备上运行 PyTorch 模型。

简介

Torch-MLU 是寒武纪（Cambricon）提供的 PyTorch 扩展库，它允许 PyTorch 模型在 MLU（机器学习单元）硬件上运行。Torch-MLU 的最大优势是完全兼容标准 PyTorch API，开发者只需少量代码修改即可将模型迁移到 MLU 设备。

主要特点

✅ 完全兼容：与标准 PyTorch API 100% 兼容
✅ 易于使用：只需几行代码即可将模型迁移到 MLU
✅ 高性能：充分利用 MLU 硬件加速能力
✅ 无缝切换：可在 CPU、GPU 和 MLU 之间轻松切换

环境要求

在运行本教程之前，请确保已安装：

Cambricon Pytorch 工具链
- 下载地址：寒武纪官网
- 安装请参考官方文档
MLU 硬件和驱动
- 寒武纪 MLU 系列加速卡
- 已安装并配置好驱动程序
Python 依赖
```
pip install torch
```

检查环境

运行以下命令检查 MLU 是否可用：

import torch
import torch_mlu

# 检查 MLU 是否可用
if torch.mlu.is_available():
    print("MLU 设备可用！")
    print(f"可用 MLU 设备数量: {torch.mlu.device_count()}")
    for i in range(torch.mlu.device_count()):
        print(f"  设备 {i}: {torch.mlu.get_device_name(i)}")
else:
    print("MLU 设备不可用，请检查 Cambricon Pytorch 工具链和驱动")

监控 MLU 设备状态

使用 cnmon 命令可以实时查看 MLU 设备的状态（类似于 NVIDIA 的 nvidia-smi）：

cnmon -smi

该命令会显示：

MLU 设备型号和计算能力
当前使用率和温度
内存使用情况
功耗信息

建议：在运行程序前、运行中和运行后查看设备状态，以监控资源使用情况和程序执行效果。这对性能调优和故障排查非常有帮助。

快速开始

1. 基础模型定义

首先，定义一个标准的 PyTorch 模型（matmul.py）：

import torch
import torch.nn as nn

class Model(nn.Module):
    """简单的矩阵乘法模型"""
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        """执行矩阵乘法 C = A @ B"""
        return torch.matmul(A, B)

def get_inputs():
    """生成随机输入"""
    N = 2048 * 2
    A = torch.rand(N, N)
    B = torch.rand(N, N)
    return [A, B]

2. 在 CPU 上运行

import torch
from matmul import Model

# 创建模型
model = Model()

# 获取输入
A, B = get_inputs()

# 执行计算（在 CPU 上）
result = model(A, B)

3. 在 MLU 上运行

import torch
import torch_mlu  # 导入 MLU 支持
from matmul import Model

# 检查 MLU 是否可用
if torch.mlu.is_available():
    print("使用 MLU 设备")
    device = 'mlu'
else:
    print("MLU 不可用，使用 CPU")
    device = 'cpu'

# 创建模型并移动到设备
model = Model().to(device)

# 获取输入并移动到设备
A, B = get_inputs()
A = A.to(device)
B = B.to(device)

# 执行计算
result = model(A, B)

# 将结果移回 CPU（如果需要）
result = result.cpu()

4. 运行完整测试

运行本教程提供的测试脚本：

python test_matmul.py

该脚本会：

在 CPU 上运行矩阵乘法
在 MLU 上运行矩阵乘法（如果可用）
对比两种设备的结果
演示 Torch-MLU 的 API 兼容性

注意：如果没有 MLU 设备，脚本会优雅地跳过 MLU 测试部分，但仍然会演示 API 兼容性。

代码详解

关键步骤

1. 导入 torch_mlu

import torch_mlu

这是使用 MLU 的第一步。导入 torch_mlu 后，PyTorch 会获得 MLU 相关的功能。

2. 检查 MLU 可用性

if not torch_mlu.is_available():
    print("MLU 设备不可用！")
    return

在实际应用中，建议先检查 MLU 是否可用，然后优雅地回退到 CPU。

3. 移动数据和模型到 MLU

# 移动模型
model = Model().to('mlu')

# 移动张量
A = A.to('mlu')
B = B.to('mlu')

使用 .to('mlu') 方法将对象移动到 MLU 设备。这与 PyTorch 中使用 .to('cuda') 的方式完全相同。

4. 同步 MLU

torch.mlu.synchronize()

在计时或需要确保计算完成时，使用 synchronize() 同步 MLU。

完整示例

import torch
import torch_mlu
import time

# 创建模型和数据
model = Model().to('mlu')
A, B = get_inputs()
A, B = A.to('mlu'), B.to('mlu')

# 预热
for _ in range(5):
    _ = model(A, B)

# 计时
torch.mlu.synchronize()
start = time.time()
for _ in range(10):
    result = model(A, B)
torch.mlu.synchronize()
elapsed = (time.time() - start) / 10 * 1000

print(f"平均执行时间: {elapsed:.3f} ms")

API 兼容性

Torch-MLU 支持大部分 PyTorch 操作。以下是一些常用操作的兼容性示例：

张量操作

import torch
import torch_mlu

# 创建张量
x = torch.randn(10, 10).to('mlu')

# 基本运算
y = x @ x.T           # 矩阵乘法
z = torch.relu(x)     # 激活函数
s = torch.sum(x)      # 归约求和
m = torch.mean(x)     # 平均值

神经网络层

import torch.nn as nn

# 线性层
linear = nn.Linear(128, 256).to('mlu')

# 卷积层
conv = nn.Conv2d(3, 64, 3).to('mlu')

# 批归一化
bn = nn.BatchNorm2d(64).to('mlu')

# 池化层
pool = nn.MaxPool2d(2).to('mlu')

损失函数

import torch.nn as nn

# 交叉熵损失
ce_loss = nn.CrossEntropyLoss().to('mlu')

# MSE 损失
mse_loss = nn.MSELoss().to('mlu')

优化器

import torch.optim as optim

# SGD 优化器
optimizer = optim.SGD(model.parameters(), lr=0.01)

# Adam 优化器
optimizer = optim.Adam(model.parameters(), lr=0.001)

训练循环示例

model = Model().to('mlu')
criterion = nn.CrossEntropyLoss().to('mlu')
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(num_epochs):
    for data, target in dataloader:
        # 移动数据到 MLU
        data, target = data.to('mlu'), target.to('mlu')
        
        # 前向传播
        output = model(data)
        loss = criterion(output, target)
        
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

常见问题

1. 如何检查 MLU 是否可用？

import torch_mlu
print(torch_mlu.is_available())

2. 如何查看 MLU 设备信息？

import torch_mlu
device_count = torch_mlu.device_count()
for i in range(device_count):
    print(f"设备 {i}: {torch_mlu.get_device_name(i)}")

3. MLU 和 CPU 结果不一致怎么办？

MLU 和 CPU 的浮点运算可能存在微小的差异，这是正常的。通常：

使用 torch.allclose() 检查结果是否在容忍范围内
对于大多数应用，rtol=1e-3, atol=1e-3 是合理的容忍度

if torch.allclose(cpu_result, mlu_result, rtol=1e-3, atol=1e-3):
    print("结果一致！")
else:
    print("结果存在差异，但在正常范围内")

4. 如何在 CPU 和 MLU 之间切换？

device = 'mlu' if torch_mlu.is_available() else 'cpu'
model = Model().to(device)
data = data.to(device)

5. MLU 内存不足怎么办？

如果遇到内存不足错误，可以：

减小批量大小（batch size）
使用梯度累积
清理不必要的中间变量
使用 torch_mlu.empty_cache() 清空缓存

# 清空 MLU 缓存
torch_mlu.empty_cache()

6. 如何调试 MLU 上的代码？

建议的开发流程：

先在 CPU 上验证：确保算法正确
再迁移到 MLU：使用 .to('mlu') 移动数据和模型
验证结果：对比 CPU 和 MLU 的输出
性能调优：根据需要进行优化

进阶话题

混合精度训练

Torch-MLU 支持自动混合精度（AMP）：

from torch_mlu.amp import autocast, GradScaler

model = Model().to('mlu')
optimizer = optim.Adam(model.parameters(), lr=0.001)
scaler = GradScaler()

for data, target in dataloader:
    data, target = data.to('mlu'), target.to('mlu')
    
    # 使用自动混合精度
    with autocast():
        output = model(data)
        loss = criterion(output, target)
    
    optimizer.zero_grad()
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

多卡训练

import torch_mlu.distributed as dist

# 初始化分布式环境
dist.init_process_group(backend='cncl')

# 创建模型并使用 DDP
model = Model().to('mlu')
model = torch.nn.parallel.DistributedDataParallel(
    model, device_ids=[torch_mlu.current_device()]
)

总结

Torch-MLU 让开发者能够轻松地将 PyTorch 模型迁移到 MLU 设备，享受硬件加速带来的性能提升。核心要点：

简单易用：导入 torch_mlu，使用 .to('mlu') 移动数据
完全兼容：与标准 PyTorch API 100% 兼容
性能优越：充分利用 MLU 硬件加速能力
灵活切换：可在不同设备间轻松切换

通过本教程，你应该已经掌握了使用 Torch-MLU 的基本方法。现在可以开始将你自己的 PyTorch 模型迁移到 MLU 上了！

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Torch-MLU 入门教程

目录

简介

主要特点

环境要求

检查环境

监控 MLU 设备状态

快速开始

1. 基础模型定义

2. 在 CPU 上运行

3. 在 MLU 上运行

4. 运行完整测试

代码详解

关键步骤

1. 导入 torch_mlu

2. 检查 MLU 可用性

3. 移动数据和模型到 MLU

4. 同步 MLU

完整示例

API 兼容性

张量操作

神经网络层

损失函数

优化器

训练循环示例

常见问题

1. 如何检查 MLU 是否可用？

2. 如何查看 MLU 设备信息？

3. MLU 和 CPU 结果不一致怎么办？

4. 如何在 CPU 和 MLU 之间切换？

5. MLU 内存不足怎么办？

6. 如何调试 MLU 上的代码？

进阶话题

混合精度训练

多卡训练

总结

参考资源

FilesExpand file tree

tutorial.md

Latest commit

History

tutorial.md

File metadata and controls

Torch-MLU 入门教程

目录

简介

主要特点

环境要求

检查环境

监控 MLU 设备状态

快速开始

1. 基础模型定义

2. 在 CPU 上运行

3. 在 MLU 上运行

4. 运行完整测试

代码详解

关键步骤

1. 导入 torch_mlu

2. 检查 MLU 可用性

3. 移动数据和模型到 MLU

4. 同步 MLU

完整示例

API 兼容性

张量操作

神经网络层

损失函数

优化器

训练循环示例

常见问题

1. 如何检查 MLU 是否可用？

2. 如何查看 MLU 设备信息？

3. MLU 和 CPU 结果不一致怎么办？

4. 如何在 CPU 和 MLU 之间切换？

5. MLU 内存不足怎么办？

6. 如何调试 MLU 上的代码？

进阶话题

混合精度训练

多卡训练

总结

参考资源