文档首页/ AI开发平台ModelArts/ ModelArts用户指南（Standard）/ 使用ModelArts Standard训练模型/ 模型训练存储加速
更新时间：2024-12-10 GMT+08:00
查看PDF
分享

模型训练存储加速

针对AI训练场景中大模型Checkpoint保存和加载带来的I/O挑战，华为云提供了基于对象存储服务OBS+高性能弹性文件服务SFS Turbo的AI云存储解决方案，如下图所示。

SFS Turbo HPC型支持和OBS数据联动，您可以通过SFS Turbo HPC型文件系统来加速对OBS对象存储中的数据访问，并将生成的结果数据异步持久化到OBS对象存储中长期低成本保存。

图1 基于OBS+SFS Turbo的存储解决方案

OBS + SFS Turbo存储加速的具体方案请查看： 面向AI场景使用OBS+SFS Turbo的存储加速实践。

设置训练存储加速

当完成上传数据至OBS并预热到SFS Turbo中步骤后，在ModelArts Standard中创建训练作业时，设置训练“SFS Turbo”，在“文件系统”中选择SFS Turbo实例名称，并指定“存储位置”和“云上挂载路径”。系统会在训练作业启动前，自动将存储位置中的文件目录挂载到训练容器中指定路径。

图2 设置训练“SFS Turbo”

当前训练作业支持挂载多个弹性文件服务SFS Turbo，文件系统支持重复挂载，但挂载路径不可重复。文件系统目录需指定已存在的目录，否则会导致训练作业异常。

然后在超参或者环境变量中设置checkpoint和数据的挂载路径。

图3 在超参或者环境变量中设置checkpoint和数据的挂载路径

训练存储加速的代码样例（PyTorch版reload ckpt）

PyTorch模型保存有两种方式。
  • 仅保存模型参数 
    state_dict = model.state_dict()
torch.save(state_dict, path)
  • 保存整个Model（不推荐） 
    torch.save(model, path)

可根据step步数、时间等周期性保存模型的训练过程的产物。

将模型训练过程中的网络权重、优化器权重、以及epoch进行保存，便于中断后继续训练恢复

checkpoint = {
"net": model.state_dict(),
"optimizer": optimizer.state_dict(),
"epoch": epoch
}
if not os.path.isdir('model_save_dir'):
os.makedirs('model_save_dir')
torch.save(checkpoint,'model_save_dir/ckpt_{}.pth'.format(str(epoch)))

以在超参中设置checkpoint和数据的挂载路径为例，完整代码示例如下：

import os
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--s3_train_url", type=str)
args, unparsed = parser.parse_known_args()
# s3_train_url 将被赋值为"/mnt/user-ckpt"
s3_train_url = args.s3_train_url

# 判断输出路径中是否有模型文件。如果无文件则默认从头训练，如果有模型文件，则加载epoch值最大的ckpt文件当做预训练模型。
if os.listdir(s3_train_url):
print('> load last ckpt and continue training!!')
last_ckpt = sorted([file for file in os.listdir(s3_train_url) if file.endswith(".pth")])[-1]
local_ckpt_file = os.path.join(s3_train_url, last_ckpt)
print('last_ckpt:', last_ckpt)
# 加载断点
checkpoint = torch.load(local_ckpt_file)
# 加载模型可学习参数
model.load_state_dict(checkpoint['net'])
# 加载优化器参数
optimizer.load_state_dict(checkpoint['optimizer'])
# 获取保存的epoch，模型会在此epoch的基础上继续训练
start_epoch = checkpoint['epoch']
start = datetime.now()
total_step = len(train_loader)
for epoch in range(start_epoch + 1, args.epochs):
for i, (images, labels) in enumerate(train_loader):
images = images.cuda(non_blocking=True)
labels = labels.cuda(non_blocking=True)
# Forward pass
outputs = model(images)
loss = criterion(outputs, labels)
# Backward and optimize
optimizer.zero_grad()
loss.backward()
optimizer.step()
...

# 保存模型训练过程中的网络权重、优化器权重、以及epoch
checkpoint = {
"net": model.state_dict(),
"optimizer": optimizer.state_dict(),
"epoch": epoch
}
if not os.path.isdir(s3_train_url):
os.makedirs(s3_train_url)
torch.save(checkpoint, os.path.join(s3_train_url, 'ckpt_best_{}.pth'.format(epoch)))

训练存储加速的代码样例（MIndSpore版reload ckpt）

import os
import argparse
from resnet import resnet50
from mindspore.nn.optim.momentum import Momentum 
from mindspore.nn.loss import SoftmaxCrossEntropyWithLogits
from mindspore import load_checkpoint, load_param_into_net
from mindspore.train import Model, CheckpointConfig, ModelCheckpoint
from mindspore.train.callback import LossMonitor

parser = argparse.ArgumentParser()
parser.add_argument("--train_url", type=str)
parser.add_argument("--batch_size", type=int, default=32, help="Batch size.") 
parser.add_argument("--num_classes", type=int, default=10, help="Num classes.") 
parser.add_argument("--do_train", type=bool, default=True, help="Do train or not.") 
args_opt, unparsed = parser.parse_known_args()
# train_url 将被赋值为"/home/ma-user/modelarts/outputs/train_url_0" 。
train_url = args_opt.train_url

# 初始定义的网络、损失函数及优化器，详细请参见MindSpore保存与加载。
# 1.初始定义的网络,以“ResNet50”为例。详细请参见ResNet50。
net = resnet50(args_opt.batch_size, args_opt.num_classes)
# 2.定义损失函数，详细请参见MindSpore自定义损失函数。
ls = SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean")
# 3.定义优化器，详细请参见MindSpore自定义优化器。
opt = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), 0.01, 0.9)
# 首次训练的epoch初始值，mindspore1.3及以后版本会支持定义epoch_size初始值。
# cur_epoch_num = 0
# 判断输出obs路径中是否有模型文件。如果无文件则默认从头训练，如果有模型文件，则加载epoch值最大的ckpt文件当做预训练模型。
if os.listdir(train_url):
    last_ckpt = sorted([file for file in os.listdir(train_url) if file.endswith(".ckpt")])[-1]
    print('last_ckpt:', last_ckpt)
    last_ckpt_file = os.path.join(train_url, last_ckpt)
     # 加载断点，详细请参见mindspore.load_checkpoint。
    param_dict = load_checkpoint(last_ckpt_file) 
    print('> load last ckpt and continue training!!')
    # 加载模型参数到net。
    load_param_into_net(net, param_dict)
    # 加载模型参数到opt。
    load_param_into_net(opt, param_dict)

    # 获取保存的epoch值，模型会在此epoch的基础上继续训练,此参数在mindspore1.3及以后版本会支持。
    # if param_dict.get("epoch_num"):
    #     cur_epoch_num = int(param_dict["epoch_num"].data.asnumpy())
model = Model(net, loss_fn=ls, optimizer=opt, metrics={'acc'})
# as for train, users could use model.train
if args_opt.do_train:
    dataset = create_dataset()
    batch_num = dataset.get_dataset_size()
    config_ck = CheckpointConfig(save_checkpoint_steps=batch_num,
                                     keep_checkpoint_max=35)
    # append_info=[{"epoch_num": cur_epoch_num}],mindspore1.3及以后版本会支持append_info参数，保存当前时刻的epoch值。
    # 保存网络参数，详细请参见mindspore.train.ModelCheckpoint。
    ckpoint_cb = ModelCheckpoint(prefix="train_resnet_cifar10",
                                     directory=args_opt.train_url,
                                     config=config_ck)
    loss_cb = LossMonitor()
    model.train(epoch_size, dataset, callbacks=[ckpoint_cb, loss_cb])
    # model.train(epoch_size-cur_epoch_num, dataset, callbacks=[ckpoint_cb, loss_cb])，mindspore1.3及以后版本支持从断点恢复训练。

相关文档