西双版纳傣族自治州网站建设_网站建设公司_在线客服_seo优化

寒武纪MLU上快速上手PyTorch指南

在AI基础设施国产化加速的今天，越来越多开发者面临从NVIDIA CUDA平台向国产芯片迁移的实际需求。寒武纪MLU凭借其高能效比和完整的软件生态，正成为信创场景下的重要选择。对于熟悉PyTorch的工程师而言，最关心的问题往往是：“我能不能像用GPU一样，直接把现有代码跑在MLU上？”

答案是肯定的——只要稍作适配。

寒武纪提供的Cambricon定制版PyTorch（基于v2.8）最大程度兼容了原生接口，只需将cuda替换为mlu，并引入torch_mlu后端支持，即可实现模型、数据与计算的无缝迁移。整个过程无需重写核心逻辑，开发体验几乎无感切换。

下面我们就从设备查看、环境配置到完整训练流程，一步步带你打通PyTorch-MLU的全链路开发路径。

设备状态监控：用cnmon替代nvidia-smi

刚拿到一台搭载MLU的服务器，第一步当然是确认硬件是否就位。这时你需要的是cnmon—— 它就是寒武纪世界的nvidia-smi。

运行最简单的命令：

cnmon

你会看到类似这样的输出：

+-----------------------------------------------------------------------------+ | CNMon Version: 1.15.0 | | Driver Version: 1.15.0 | +-------------------------------+---------------------+-----------------------+ | Card ID | Temp (C) | Memory | Used / Total (MiB) | Process | +---------+----------+----------+--------------------+-----------------------+ | 0 | 48 | 32768 | 1024 / 32768 | python(2345) | +---------+----------+----------+--------------------+-----------------------+

这说明你的系统已识别出一块MLU卡，驱动正常加载，显存可用。

如果想盯住某一张卡看资源变化，可以用自动刷新模式：

cnmon -r -t 1000

每秒更新一次，特别适合在训练过程中观察显存占用和温度波动。

也可以指定查看特定设备：

cnmon -c 0 # 查看第0号卡

这个工具虽然界面朴素，但足以覆盖日常调试所需的基本信息。更详细的参数可参考官方文档，但我们通常只需要知道“有没有”、“够不够”、“热不热”这三个关键点就够了。

环境准备：导入torch_mlu的正确姿势

如果你使用的是官方预置镜像（比如Docker容器或虚拟机），大概率已经集成了以下组件：

• Python 3.9+
• Cambricon定制版 PyTorch v2.8
• torch_mlu插件
• CNNL、CNCL、MagicMind Runtime

这意味着你不需要手动编译或安装任何底层库，开箱即用。

但有一个非常关键的细节：导入顺序不能错。

import torch # 必须先导入标准PyTorch import torch_mlu # 再激活MLU后端

为什么？因为torch_mlu实际上是一个设备后端插件，它会在初始化时注册'mlu'设备类型到全局上下文中。如果torch尚未加载，这个注册机制就会失效，导致后续调用torch.mlu.is_available()返回False，甚至出现设备不可见的问题。

验证一下是否一切就绪：

print(torch.mlu.is_available()) # 应该输出 True print(torch.mlu.device_count()) # 输出当前可见的MLU数量，如 1 或 2

一旦返回True，恭喜你，已经成功打通了通往MLU世界的大门。

接下来可以试试在MLU上创建张量：

x = torch.randn(3, 3).to('mlu') y = torch.ones_like(x).mlu() # 另一种写法 z = x + y print(z.device) # 输出: mlu:0

你会发现语法完全一致，唯一的区别只是把'cuda'换成了'mlu'。这种设计极大降低了迁移成本，老用户几乎零学习曲线。

模型迁移：.to('mlu')全家桶

模型部署到MLU的方式与CUDA一模一样，核心就是一句话：

model = model.to('mlu')

无论是自己定义的网络，还是torchvision.models里的经典结构，都可以通过这种方式整体搬移。

例如加载ResNet18：

import torchvision.models as models model = models.resnet18(pretrained=False) model = model.to('mlu') # 所有权重视图迁移到MLU

此时模型的所有参数都位于MLU显存中，前向传播和反向传播会自动调度至MLU执行，无需额外干预。

多卡训练怎么搞？

如果你的机器装了多块MLU卡，可以通过DataParallel实现单机多卡并行：

if torch.mlu.device_count() > 1: model = nn.DataParallel(model, device_ids=[0, 1]) model = model.to('mlu')

注意目前不支持 DistributedDataParallel（DDP），所以跨节点分布式训练暂时无法实现。但对于大多数中小规模任务来说，单机双卡或四卡已经足够。

一个实用建议：在启动脚本开头加个日志打印，避免误判设备状态：

device = torch.device('mlu' if torch.mlu.is_available() else 'cpu') print(f"Running on device: {device}")

这样即使环境异常也能第一时间发现问题。

损失函数别忘了也得“上车”

新手常踩的一个坑是：模型上了MLU，数据也搬过去了，结果训练时报错：

RuntimeError: Expected all tensors to be on the same device

原因往往出在损失函数身上。

默认情况下，nn.Module创建的损失层（如CrossEntropyLoss）仍然驻留在CPU上，必须显式移动：

criterion = nn.CrossEntropyLoss() criterion = criterion.to('mlu') # 关键一步！

或者统一用设备变量控制：

device = torch.device('mlu') criterion = criterion.to(device)

同理，所有参与计算的模块——包括自定义的Metric、Loss、Head等——都要确保与输入数据在同一设备上。

一个小技巧：可以在模型.to()之后顺手把整个模型相关的组件一起迁移：

model = model.to(device) criterion = criterion.to(device)

保持一致性，才能避免“设备错配”的 runtime 错误。

数据搬运策略：效率与简洁之间的平衡

数据默认从 DataLoader 加载出来是在CPU上的，需要主动拷贝到MLU设备。

常见做法有两种：

方式一：迭代中即时搬运（推荐）

for data, target in train_loader: data = data.to('mlu', non_blocking=True) target = target.to('mlu', non_blocking=True) output = model(data) loss = criterion(output, target) ...

其中non_blocking=True表示启用异步传输（若硬件支持DMA），可以让数据搬运和前一轮计算重叠，提升流水线效率。

不过要注意，并非所有场景都适合开启non_blocking。在小批量或低延迟推理中反而可能增加调度开销，建议在大批量训练中启用。

方式二：使用.mlu()快捷方法

data = data.mlu() target = target.mlu()

这是torch_mlu提供的便捷封装，语义清晰且代码更短，适合交互式调试或Notebook场景。

两种方式本质相同，选哪个取决于你的编码风格和性能要求。

完整训练示例：MNIST分类实战

理论讲完，来点真家伙。下面是一个可在MLU上运行的完整MNIST训练脚本，涵盖数据加载、模型构建、训练循环与评估全流程。

import torch import torch_mlu import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms from torch.optim.lr_scheduler import StepLR import torch.nn.functional as F class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(1, 32, 3, 1) self.conv2 = nn.Conv2d(32, 64, 3, 1) self.dropout1 = nn.Dropout2d(p=0.25) self.dropout2 = nn.Dropout2d(p=0.5) self.fc1 = nn.Linear(9216, 128) self.fc2 = nn.Linear(128, 10) def forward(self, x): x = self.conv1(x) x = F.relu(x) x = self.conv2(x) x = F.relu(x) x = F.max_pool2d(x, 2) x = self.dropout1(x) x = torch.flatten(x, 1) x = self.fc1(x) x = F.relu(x) x = self.dropout2(x) x = self.fc2(x) return F.log_softmax(x, dim=1) def train(model, device, train_loader, optimizer, epoch): model.train() for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) optimizer.zero_grad() output = model(data) loss = F.nll_loss(output, target) loss.backward() optimizer.step() if batch_idx % 100 == 0: print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( epoch, batch_idx * len(data), len(train_loader.dataset), 100. * batch_idx / len(train_loader), loss.item())) def test(model, device, test_loader): model.eval() test_loss = 0 correct = 0 with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) output = model(data) test_loss += F.nll_loss(output, target, reduction='sum').item() pred = output.argmax(dim=1, keepdim=True) correct += pred.eq(target.view_as(pred)).sum().item() test_loss /= len(test_loader.dataset) accuracy = 100. * correct / len(test_loader.dataset) print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format( test_loss, correct, len(test_loader.dataset), accuracy)) def main(): device = torch.device('mlu' if torch.mlu.is_available() else 'cpu') print(f"Using device: {device}") transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_dataset = datasets.MNIST('../data', train=True, download=True, transform=transform) test_dataset = datasets.MNIST('../data', train=False, transform=transform) train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False) model = Net().to(device) optimizer = optim.Adadelta(model.parameters(), lr=1.0) scheduler = StepLR(optimizer, step_size=1, gamma=0.7) epochs = 5 for epoch in range(1, epochs + 1): train(model, device, train_loader, optimizer, epoch) test(model, device, test_loader) scheduler.step() torch.save(model.state_dict(), "mnist_cnn_mlu.pt") print("Model saved to mnist_cnn_mlu.pt") if __name__ == '__main__': main()

这段代码可以直接保存为mnist.py并运行：

python mnist.py

不出意外的话，你应该能看到类似这样的输出：

Using device: mlu:0 Train Epoch: 1 [0/60000 (0%)] Loss: 2.312345 ... Test set: Average loss: 0.0781, Accuracy: 9768/10000 (98%)

模型最终也会被保存为.pt文件，可用于后续推理或其他设备加载。

💡 提示：第一次运行时若提示找不到数据目录，可手动创建../data文件夹，或修改路径为绝对路径。

开发方式选择：Jupyter 还是 SSH？

不同开发阶段适合不同的接入方式。

Jupyter Notebook：交互式调试利器

对于算法探索、可视化分析或教学演示，Jupyter 是首选。

启动服务后，在浏览器中打开Notebook，可以直接运行如下代码：

import torch import torch_mlu print(torch.mlu.is_available()) # 输出 True x = torch.randn(2, 2).mlu() print(x)

实时反馈 + 分步调试 + 图形展示，非常适合快速验证想法。

✅ 优势：直观、灵活、易分享
❌ 缺点：不适合长时间任务，资源管理较弱

SSH终端：生产级开发标配

真正做项目开发、跑长周期训练任务，还是得靠SSH登录服务器，走命令行路线。

ssh user@mlu-server-ip -p 22

连接成功后，你可以：

• 编辑脚本（vim/nano）
• 提交批量任务（bash/python）
• 监控进程（top/cnmon）
• 管理文件（scp/rsync）

配合tmux或screen工具还能实现后台持久化运行：

tmux new-session -d -s train 'python mnist.py'

即使本地断网，训练也不会中断。

✅ 优势：稳定、高效、可自动化
❌ 缺点：门槛略高，需熟悉Linux操作

建议组合使用：前期用Jupyter快速原型，后期用SSH部署训练。

这种高度集成的设计思路，正推动着国产AI算力向更可靠、更高效的开发体验迈进。随着生态不断完善，未来我们或许不再需要强调“迁移”本身——因为它本就该是透明的。