幻方萤火AI算力平台使用笔记（3）：Workspace功能浅试——Resnet-50训练示例

携手创作，共同成长！这是我参与「掘金日新计划 · 8 月更文挑战」的第5天，点击查看活动详情

导语

前面的两篇博客幻方萤火AI算力平台使用笔记（1）：MacOS下的客户端安装与管理平台概览和幻方萤火AI算力平台使用笔记（2）：Hfai本地命令行工具使用，我们介绍了如何配置MacOS端的网络vpn和本地hfai命令行工具，在完成以上基础配置后，本篇博客将尝试使用算力平台运行一个小的example作为示例，以明确开发流程，本文将以幻方官方文档提供的Resnet-50代码作为示例。

hfai workspace简介

幻方的萤火超算是一个远程的服务集群。对于研究者来说，他们需要将本地的数据、代码传到这个远程集群中，才可以训练其模型。hfai workspace就是用来连通本地工程目录和远程萤火集群的工具。

通过该工具，用户可以在自己环境，如个人电脑、个人集群等，编写代码，调试模型，然后通过一行命令直接将调试好的代码上传远程萤火集群，利用超算算力来训练模型。

一个简单的示例：Resnet-50网络训练

为了更加直观的理解hfai workspace的功能，这里介绍一个快速上手的案例进行运行。

Resnet-50简介

Resnet是是经典的残差网络(Residual Network)的缩写，该系列网络广泛用于目标分类等领域以及作为计算机视觉任务主干经典神经网络的一部分，其中一些典型的网络有Resnet-50, Resnet-101等。Resnet网络在计算机视觉领域应用广泛，是一个非常经典的模型。

Resnet-50训练代码

本篇博客中使用的训练文件train.py代码来自幻方萤火平台文档快速开始中推荐的模型库中的模型，具体如下：

import hf_env
hf_env.set_env('202111')

import os
from pathlib import Path
import torch
from torch import nn
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data.distributed import DistributedSampler
from torch.optim import SGD
from torch.optim.lr_scheduler import StepLR
from torchvision import transforms, models

import hfai
import hfai.distributed as dist


def train(dataloader, model, criterion, optimizer, epoch, local_rank, start_step, best_acc):
    model.train()

    for step, batch in enumerate(dataloader):
        step += start_step

        samples, labels = [x.cuda(non_blocking=True) for x in batch]
        outputs = model(samples)
        optimizer.zero_grad()
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        if local_rank == 0 and step % 20 == 0:
            print(f'Epoch: {epoch}, Step: {step}, Loss: {loss.item()}', flush=True)

        # 保存
        model.try_save(epoch, step + 1, others=best_acc)


def validate(dataloader, model, criterion, epoch, local_rank):
    loss, correct1, correct5, total = torch.zeros(4).cuda()
    model.eval()
    with torch.no_grad():
        for step, batch in enumerate(dataloader):
            samples, labels = [x.cuda(non_blocking=True) for x in batch]
            outputs = model(samples)
            loss += criterion(outputs, labels)
            _, preds = outputs.topk(5, -1, True, True)
            correct1 += torch.eq(preds[:, :1], labels.unsqueeze(1)).sum()
            correct5 += torch.eq(preds, labels.unsqueeze(1)).sum()
            total += samples.size(0)

    for x in [loss, correct1, correct5, total]:
        dist.reduce(x, 0)

    if local_rank == 0:
        loss_val = loss.item() / dist.get_world_size() / len(dataloader)
        acc1 = 100 * correct1.item() / total.item()
        acc5 = 100 * correct5.item() / total.item()
        print(f'Epoch: {epoch}, Loss: {loss_val}, Acc1: {acc1:.2f}%, Acc5: {acc5:.2f}%', flush=True)

    return correct1.item() / total.item()


def main(local_rank):
    # 超参数设置
    epochs = 2
    batch_size = 50
    lr = 0.1
    save_path = 'output/resnet'
    Path(save_path).mkdir(exist_ok=True, parents=True)

    # 多机通信
    ip = os.environ['MASTER_ADDR']
    port = os.environ['MASTER_PORT']
    hosts = int(os.environ['WORLD_SIZE'])  # 机器个数
    rank = int(os.environ['RANK'])  # 当前机器编号
    gpus = torch.cuda.device_count()  # 每台机器的GPU个数

    # world_size是全局GPU个数，rank是当前GPU全局编号
    dist.init_process_group(backend='nccl',
                            init_method=f'tcp://{ip}:{port}',
                            world_size=hosts * gpus,
                            rank=rank * gpus + local_rank)
    torch.cuda.set_device(local_rank)

    # 模型、数据、优化器
    model = models.resnet50()
    model = DistributedDataParallel(model.cuda(), device_ids=[local_rank])

    train_transform = transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])
    ])  # 定义训练集变换
    train_dataset = hfai.datasets.ImageNet('train', transform=train_transform)
    train_datasampler = DistributedSampler(train_dataset)
    train_dataloader = train_dataset.loader(batch_size, sampler=train_datasampler, num_workers=4, pin_memory=True)

    val_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])
    ])  # 定义测试集变换
    val_dataset = hfai.datasets.ImageNet('val', transform=val_transform)
    val_datasampler = DistributedSampler(val_dataset)
    val_dataloader = val_dataset.loader(batch_size, sampler=val_datasampler, num_workers=4, pin_memory=True)

    criterion = nn.CrossEntropyLoss()
    optimizer = SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)
    scheduler = StepLR(optimizer, step_size=30, gamma=0.1)

    ckpt_path = os.path.join(save_path, 'latest.pt')
    start_epoch, start_step, best_acc = hfai.checkpoint.init(model, optimizer, scheduler=scheduler, ckpt_path=ckpt_path)
    best_acc = best_acc or 0

    # 训练、验证
    for epoch in range(start_epoch, epochs):
        # resume from epoch and step
        train_datasampler.set_epoch(epoch)
        train_dataloader.set_step(start_step)

        train(train_dataloader, model, criterion, optimizer, epoch, local_rank, start_step, best_acc)
        start_step = 0  # reset
        scheduler.step()
        acc = validate(val_dataloader, model, criterion, epoch, local_rank)
        # 保存
        if rank == 0 and local_rank == 0:
            if acc > best_acc:
                best_acc = acc
                print(f'New Best Acc: {100*acc:.2f}%!')
                torch.save(model.module.state_dict(),
                           os.path.join(save_path, 'best.pt'))


if __name__ == '__main__':
    ngpus = torch.cuda.device_count()
    torch.multiprocessing.spawn(main, args=(), nprocs=ngpus)

该训练代码会在ImageNet数据集上对Resnet-50进行epochs = 2的训练（这里是为了快速演示，所以设置了一个较小的epochs）。

该训练代码使用分布式数据并行（DistributedDataParallel，即DDP）训练，并在每个epoch训练完成后进行一次验证，统计Top-1和Top-5的分类准确率。最后，会将最好的模型checkpoint进行存储，保存为best.pt。

最后，需要注意的是我们在文件开头加入了

import hf_env 
hf_env.set_env('202111')

设置该训练文件的运行环境，这里直接使用了幻方官方的2021年11月的运行环境。

hfai workspace初始化

假设我们在本地已经调试开发好了一个在ImageNet上进行训练的Resnet-50模型的训练代码，接下来，我们需要将代码上传到算力平台进行训练。

首先，我们需要初始化该项目的workspace，即该模型对应的项目文件夹。首先，我们需要手动创建一个本地的文件夹，这里我取名为workspace_example_for_resnet

(hfai) jiexing@jiexingdeMacBook-Pro hfai % mkdir workspace_example_for_resnet 
(hfai) jiexing@jiexingdeMacBook-Pro hfai % cd workspace_example_for_resnet 

接着，我们将训练的文件加入到改文件夹中，训练文件命名为train.py，并添加了一个Readme文件。

(hfai) jiexing@jiexingdeMacBook-Pro workspace_example_for_resnet % ls
README.md       train.py

之后，我们进行初始化操作，命令为hfai workspace init [workspace_name]，输出如下：

(hfai) jiexing@jiexingdeMacBook-Pro workspace_example_for_resnet % hfai workspace init workspace_example_for_resnet
查询用户 group 信息...
初始化 workspace [/Users/jiexing/Nextcloud/project/hfai/workspace_example_for_resnet]->[oss://group_linzhh/linzhh/workspaces/workspace_example_for_resnet] 成功

hfai workspace推送与执行

初始化完成后，我们将这两个文件推送到远程平台端，使用命令hfai workspace push

(hfai) jiexing@jiexingdeMacBook-Pro workspace_example_for_resnet % hfai workspace push
开始遍历本地工作区目录...
开始遍历集群工作区目录...
开始打包本地工作区目录...
(1/2) 开始同步本地目录 /Users/jiexing/Nextcloud/project/hfai/workspace_example_for_resnet 到远端，共3.05KB...
pushing ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:00
(2/2) 上传成功，开始同步到集群，请等待...
syncing ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:00
推送成功

hfai工具会自动将工作目录进行打包和上传，推送至幻方萤火平台。而对于执行，hfai提供了多重模式的执行方式，具体可以使用hfai python -h查询。

总体来说，hfai提供三种模式的执行方式：

1. 本地执行。这种方式和普通的python脚本在本地环境中执行没有区别，使用的命令为hfai python <experiment.py> [<experiment_params>...]；
2. 远程执行。这种方式即提交任务到萤火二号运行，使用的命令为hfai python <experiment.py> [<experiment_params>...] -- [CLUSTER_OPTIONS]，相比于1，这里使用--命令添加后续远程参数；
3. 模拟执行。这种方式在本地模拟运行，会提供萤火二号一致的环境变量和表现，主要测试本地编写的代码脚本能否收到幻方管理平台的任务打断信号等，使用的命令为hfai python <experiment.py> [<experiment_params>...] ++ [CLUSTER_OPTIONS]，相比于1，这里使用++命令添加模拟参数。

这里，我们使用官方仓库推荐的运行命令

hfai python train.py -- -n 1

该命令表示我们要使用远程算力平台运行train.py，使用的服务器节点数为1（一台机器8张卡）。

运行过程如下：

(hfai) jiexing@jiexingdeMacBook-Pro workspace_example_for_resnet % hfai python train.py -- -n 1
检测到是 [/Users/jiexing/Nextcloud/project/hfai/workspace_example_for_resnet] -> [group_linzhh/linzhh/workspaces/workspace_example_for_resnet] 中的代码，先推送到远端...
开始遍历本地工作区目录...
开始遍历集群工作区目录...
数据已同步，忽略本次操作
……
====================  fetching  log on rank 0... ====================
[2022-08-01 09:17:40.756442] [start training train.py on jd-h0605-dl for linzhh]
[2022-08-01 09:17:40.855583] [训练前检查] 检查[jd-h0605-dl] cpu memory 44.0 < 150G and total gpu memory 0 < 100M 通过
[2022-08-01 09:18:10.096679] Epoch: 0, Step: 0, Loss: 7.289088249206543
[2022-08-01 09:18:13.314214] Epoch: 0, Step: 20, Loss: 7.264989852905273
[2022-08-01 09:18:15.124726] Epoch: 0, Step: 40, Loss: 7.147526264190674
……
[2022-08-01 09:28:06.848568] Epoch: 1, Step: 3200, Loss: 3.582792282104492
[2022-08-01 09:28:15.918022] Epoch: 1, Loss: 3.25835107421875, Acc1: 30.57%, Acc5: 56.37%
[2022-08-01 09:28:17.562394] New Best Acc: 30.57%!
[2022-08-01 09:28:21.328877] [finish training train.py on jd-h0605-dl for linzhh]

输出查看与下载

我们可以登录网页端的管理平台（studio.yinghuo.high-flyer.cn/ ），在工作区栏目中找到我们这次运行使用的工作区目录，即可查看这次的输出

我们可以通过workspace的同步命令，将这些内容下载到本地，这些内容将在后面的博客中记录。

总结

本篇博客为大家介绍了幻方萤火AI算力平台上一个最简单的demo示例程序的运行，我们通过在本地编写Resnet-50模型的代码，将其推送到远程后在Imagenet数据集上进行训练，得到了预期的输出结果。下篇博客将记录workspace相关命令的学习笔记，敬请期待。