Zero-to-GPU / 零基础上云训练

从安装 Python 包到训练 MNIST MLP

这份教程面向第一次使用 DuanFlow 的老师和学生。你会看到一份普通 PyTorch MNIST MLP 训练代码,其中 黄色高亮 的几行就是 DuanFlow 相关代码;其它部分和在本地电脑或学校机房里跑 PyTorch 没区别。

核心思路

普通代码照旧写 模型、数据集、训练循环、loss、optimizer 都是标准 PyTorch。老师原来的教学代码不用重写。
只加高亮几行 高亮的 DuanFlow 代码负责声明镜像、选择 GPU/NPU,并把函数提交到服务器运行。
看代码时只记住一句话:

没有高亮的部分 = 本地 PyTorch 训练;黄色高亮的部分 = 把这段训练放到 DuanFlow 算力服务器上跑。

1. 安装 DuanFlow Python 包

在自己的电脑、学校实验室电脑或教师机上执行。这里安装的是提交任务的 SDK,不是在本机训练大模型。

terminal
python -m pip install https://duanflow.zeabur.app/sdk/duanflow-0.1.0.tar.gz

2. 浏览器授权登录

学校可以给每个老师、学生开账号。登录后,SDK 会带着你的身份提交任务,后台就能统计谁用了多少小时、产生多少费用。不要把密码、API key 或 token 粘贴给 Agent。

terminal
duanflow auth login --api-base https://duanflow.zeabur.app
duanflow whoami

3. 范例 A:使用昇腾 Ascend 910B 训练 MNIST MLP

下面代码中,黄色高亮的行是 DuanFlow 相关代码;其余就是标准 PyTorch 训练代码。昇腾 910B 使用 torch-npu,所以镜像依赖里写了 torchtorch-npu

mnist_ascend_910b.py
import duanflow as df  # DuanFlow:引入平台 SDK

image = df.Image.from_registry("registry.cmcc.example.com/synapse/ascend-runtime:latest").uv_pip_install("torch", "torch-npu", "numpy")  # DuanFlow:声明服务器镜像和依赖
app = df.App("school-mnist-ascend", image=image)  # DuanFlow:创建一个云端应用

@app.function(gpu="Ascend 910B", gpu_count=1, timeout=1800)  # DuanFlow:指定使用 1 张昇腾 910B
def train_mnist():
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    from torch.utils.data import DataLoader
    from torchvision import datasets, transforms

    import torch_npu  # 昇腾环境需要;如果在本地 CPU 跑,可以删除这一行
    device = torch.device("npu:0")  # 普通 PyTorch 写法,只是设备从 cuda 换成 npu

    # 下载并标准化 MNIST。这里和本地训练完全一样。
    tfm = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
    train_ds = datasets.MNIST("/tmp/mnist", train=True, download=True, transform=tfm)
    test_ds = datasets.MNIST("/tmp/mnist", train=False, download=True, transform=tfm)
    train_loader = DataLoader(train_ds, batch_size=128, shuffle=True)
    test_loader = DataLoader(test_ds, batch_size=1024)

    # 一个最简单的 MLP:784 -> 256 -> 10
    model = nn.Sequential(nn.Flatten(), nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10)).to(device)
    opt = torch.optim.AdamW(model.parameters(), lr=1e-3)

    # 训练 3 个 epoch。这里仍然是普通 PyTorch。
    for epoch in range(3):
        model.train()
        for x, y in train_loader:
            x, y = x.to(device), y.to(device)
            opt.zero_grad()
            loss = F.cross_entropy(model(x), y)
            loss.backward()
            opt.step()
        print(f"epoch={epoch + 1} loss={loss.item():.4f}")

    # 测试准确率。
    model.eval()
    correct = total = 0
    with torch.no_grad():
        for x, y in test_loader:
            pred = model(x.to(device)).argmax(1).cpu()
            correct += (pred == y).sum().item()
            total += y.numel()
    return {"accuracy": correct / total}

print(train_mnist.remote())  # DuanFlow:把函数提交到服务器运行;本地只负责发起任务

4. 范例 B:使用 NVIDIA RTX 4090 训练 MNIST MLP

如果学校或平台接入的是 NVIDIA 4090,代码几乎一样,只是镜像和 GPU 名称不同,设备使用 cuda

mnist_4090.py
import duanflow as df  # DuanFlow:引入平台 SDK

image = df.Image.debian_slim().uv_pip_install("torch", "torchvision", "numpy")  # DuanFlow:声明 Python 依赖
app = df.App("school-mnist-4090", image=image)  # DuanFlow:创建一个云端应用

@app.function(gpu="RTX 4090", gpu_count=1, timeout=1800)  # DuanFlow:指定使用 1 张 NVIDIA 4090
def train_mnist():
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    from torch.utils.data import DataLoader
    from torchvision import datasets, transforms

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 本地 PyTorch 常见写法

    tfm = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
    train_ds = datasets.MNIST("/tmp/mnist", train=True, download=True, transform=tfm)
    test_ds = datasets.MNIST("/tmp/mnist", train=False, download=True, transform=tfm)
    train_loader = DataLoader(train_ds, batch_size=128, shuffle=True)
    test_loader = DataLoader(test_ds, batch_size=1024)

    model = nn.Sequential(nn.Flatten(), nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10)).to(device)
    opt = torch.optim.AdamW(model.parameters(), lr=1e-3)

    for epoch in range(3):
        model.train()
        for x, y in train_loader:
            x, y = x.to(device), y.to(device)
            opt.zero_grad()
            loss = F.cross_entropy(model(x), y)
            loss.backward()
            opt.step()
        print(f"epoch={epoch + 1} loss={loss.item():.4f}")

    model.eval()
    correct = total = 0
    with torch.no_grad():
        for x, y in test_loader:
            pred = model(x.to(device)).argmax(1).cpu()
            correct += (pred == y).sum().item()
            total += y.numel()
    return {"accuracy": correct / total}

print(train_mnist.remote())  # DuanFlow:提交到服务器;结果和日志回到本地

5. 运行与查看日志

运行方式和普通 Python 文件一样。区别是:训练不会占用本地电脑 GPU,而是提交到学校后台配置的 DuanFlow 算力服务器。

terminal
python mnist_ascend_910b.py

# 或者
python mnist_4090.py
学生看到什么终端会打印训练日志和最终 accuracy。学校后台会记录这次任务。
老师看到什么老师能在学校后台看到谁跑了任务、用了多少卡时、是否成功、日志是什么。