本文介绍如何在 PyTorch 任务中访问和存储 TOS 数据。

前提条件

在使用 TOS 之前，请确保您已注册火山引擎账号并完成实名认证。

请确保您已开通 TOS 服务并创建 TOS 存储桶，具体步骤，请参见快速入门。

必须使用 Linux x86-64 操作系统。

已安装 Python 3.8 ~ 3.11 版本。

环境准备

对接 TOS 前，请先在 Python 3.8 ~ 3.11 环境安装 PyTorch。更多信息，请参见 PyTorch。

# 安装 torch 基本依赖
pip3 install torch torchvision

# 安装 torch lighting
pip3 install lightning

# 安装 torchdata
pip3 install torchdata==0.9.0

安装完成后，您可以通过以下示例验证是否安装成功。

import torch
import torchvision

# 将 checkpoint 保存为本地文件
if __name__ == '__main__':
model = torchvision.models.resnet18()

with open('test.pt', 'wb') as writer:
torch.save(model.state_dict(), writer)

with open('test.pt', 'rb') as reader:
stat_dict = torch.load(reader)

model.load_state_dict(stat_dict)
print(model)

对接方式

对接方式	说明
使用 tosfsspec 对接	tosfsspec 是火山引擎基于 fsspec 规范封装的 TOS 文件系统接口，用于将 TOS 桶对象封装为 Python 标准文件读写接口，可直接与 torchdata、PyTorch 等训练框架集成。更多信息，请参见 tosfsspec 概述。
使用 FSX 对接	FSX 支持将存储桶挂载到 Linux 系统作为本地文件系统使用，适用于需要按照文件系统语义使用存储桶中数据的场景，例如 AI 机器学习场景下需要按照文件系统语义训练模型数据。 FSX 客户端处于邀测状态，如需试用，请提交工单。支持的运行环境及约束限制详情，请参见 FSX 概述。
使用 S3 Connector 对接	S3 Connector 是 AWS 提供的 PyTorch 连接器，用于增强 torchdata 提供的 S3 相关访问能力，支持 map-style datasets、iterable-style datasets 以及 checkpointing 接口。更多信息，请参见 S3-connector-for-PyTorch。当前仅支持 TOS FNS 桶，即扁平桶。更多信息，请参见存储桶类型。

操作步骤

不同对接方式的操作步骤如下。

使用 tosfsspec 对接

使用 FSX 对接

使用 S3 Connector 对接

执行以下命令安装 tosfsspec。pip 会自动安装所需依赖（fsspec ≥ 2023.5.0、tos ≥ 2.8.0、tosnativeclient ≥ 1.1.13），无需手动安装。

# 使用 pip 直接安装 tosfsspec。
pip install tosfsspec

获取凭证信息。

地域（Region）：您的 TOS 存储桶所在的地域，例如 cn-beijing。

访问密钥（Access Key）：包含 Access Key ID (AK) 和 Secret Access Key (SK)，是访问火山引擎 API 的身份凭证。

访问域名（Endpoint）：可选参数。如果留空，tosfsspec 会根据您提供的 region 自动生成。

说明

关于地域和访问域名的详细列表，请参见地域和访问域名（Endpoint）。

关于如何创建和管理访问密钥，请参见 API 访问密钥（Access Key）。

配置长期访问凭证。

初始化 SDK 前，您需要先配置长期访问凭证，操作步骤如下：

# 打开终端并执行以下命令打开文件。
sudo vim /etc/profile

# 在文件末尾添加 AKSK 信息。
export TOS_ACCESS_KEY=AKTPYmI1Z****
export TOS_SECRET_KEY=T1dJM01UU****

# 保存文件并退出,然后执行以下命令生效配置信息。
source /etc/profile

# 验证配置信息。
echo $TOS_ACCESS_KEY
echo $TOS_SECRET_KEY

执行以下命令初始化客户端。

import tosfsspec
import os

# 配置访问密钥等信息
if __name__ == '__main__':
fs = tosfsspec.TosFileSystem(
key=os.getenv('TOS_ACCESS_KEY'),
secret=os.getenv('TOS_SECRET_KEY'),
endpoint='your endpoint',
region='your endpoint',
)

运行以下代码，使用 iterable-style datasets。

本文以使用北京地域的 bucketname 桶为例，建议您根据业务需求修改代码。

from torchdata.datapipes.iter import FSSpecFileLister, FSSpecFileOpener
import os

if __name__ == '__main__':
kwargs = {
'key': os.getenv('TOS_ACCESS_KEY'),
'secret': os.getenv('TOS_SECRET_KEY'),
'endpoint':'your endpoint',
'region':'your region'
}
bucket = 'bucketname'
print('测试 iterable-style datasets')
file_lister = FSSpecFileLister(root='tos://' + bucket + '/shard-data/', **kwargs)
iterable_dataset = FSSpecFileOpener(file_lister, mode='rb', **kwargs)
# 遍历并读取 Dataset 中的对象
for _, item in iterable_dataset:
print(item.bucket, item.key)
data = item.read()
print(len(data))

运行以下代码，保存或加载 checkpoint。

本文以使用北京地域的 bucketname 桶为例，建议您根据业务需求修改代码。

import torch
import torchvision
import tosfsspec
import TosFileSystem
import os

if __name__ == '__main__':
fs = tosfsspec.TosFileSystem(
key=os.getenv('TOS_ACCESS_KEY'),
secret=os.getenv('TOS_SECRET_KEY'),
endpoint='your endpoint',
region='your region',
)
bucket = 'bucketname'
print('测试 checkpointing')
model = torchvision.models.resnet18()

# 将 TOS 封装为 Python 标准读写接口
with fs.open(f'tos://{bucket}/checkpoint/epoch1.ckpt', 'wb') as writer:
torch.save(model.state_dict(), writer)

with fs.open(f'tos://{bucket}/checkpoint/epoch1.ckpt', 'rb') as reader:
stat_dict = torch.load(reader)
model.load_state_dict(stat_dict)
print(model)

本文以 veLinux 环境为例，介绍 FSX 客户端安装步骤及挂载整个存储桶的步骤，更多环境及参数详情，请参见以下文档：

下载与安装 FSX

FSX 配置与挂载

FSX 挂载参数说明

安装 FSX 客户端。

根据操作系统类型，请提交工单联系技术团队，获取最新的安装包。

执行以下命令，安装客户端。

dpkg -i fsx_amd64.deb

执行以下命令，查看客户端版本。

fsx -v

返回如下版本信息，表示安装成功。

Version: 1.1.0.400

配置挂载权限。

IAM 用户挂载存储桶时，需要支持 ListBucket和HeadBucket 权限，对挂载的桶内对象进行读操作时需要支持 GetObject 权限，写入文件需要支持 PutObject 和 AbortMultipartUpload 权限。

以下是一个最小权限示例，更多信息，请参见权限配置。

{
"Statement": [{
"Effect": "Allow",
"Action": [
"tos:GetObject",
"tos:PutObject",
"tos:DeleteObject",
"tos:AbortMultipartUpload",
"tos:ListBucket",
"tos:HeadBucket"
],
"Resource": [
"*"
]
}]
}

挂载 TOS 存储桶。

挂载参数详情，请参见 FSX 配置与挂载。

进入 Linux 操作系统，执行以下命令，创建挂载点。

mkdir /mnt/point

执行以下命令，挂载存储桶。

mount -t fsx /tos/<挂载的桶名> <挂载点> -o region="<地域>",endpoint="<访问域名>",access_key_id="AK",secret_access_key="SK",no_writeback_cache

示例如下：

mount -t fsx /tos/bucketname /mnt/point -o region="cn-beijing",endpoint="http://tos-cn-beijing.volces.com",access_key_id="AKAKAK****",secret_access_key="SKSKSK****",no_writeback_cache

执行以下命令，查看存储桶挂载情况。

df -h

在返回内容中看到挂载点、挂载桶名，说明挂载成功。

Filesystem Size Used Avail Use% Mounted on
BFA281845F:tos-fuse-test 8.0E 0 8.0E 0% /mnt/fuse_shm

运行以下代码，构建 iterable-style datasets。

说明

本文以使用 shard-data/ 目录为例，建议您根据业务需求修改代码。

from torchdata.datapipes.iter import FileLister, FileOpener

if __name__ == '__main__':
dataset_uri = './root/shard-data/'

print('测试 iterable-style datasets')
file_lister = FileLister(root=dataset_uri)
iterable_dataset = FileOpener(file_lister, mode='rb')
# 遍历并读取Dataset中的对象。
for _, item in iterable_dataset:
print(item.bucket, item.key)
data = item.read()
print(len(data))

运行以下代码，保存或加载 checkpoint。

说明

本文以使用 checkpoint/ 目录为例，建议您根据业务需求修改代码。

import torch
import torchvision

if __name__ == '__main__':
print('测试 checkpointing')
model = torchvision.models.resnet18()

# 直接将 tos 作为本地文件处理
with open('./root/checkpoint/epoch0.ckpt', 'wb') as writer:
torch.save(model.state_dict(), writer)

with open('./root/checkpoint/epoch0.ckpt', 'rb') as reader:
stat_dict = torch.load(reader)

model.load_state_dict(stat_dict)
print(model)

安装 s3 connector。

执行以下命令安装 s3 connector 基础库。

pip install s3torchconnector

（可选）执行以下命令安装 s3 connector for lighting。

pip install s3torchconnector[lightning]

配置环境。

执行以下命令编辑密钥配置文件。

vim ~/.aws/credentials

执行以下命令配置访问密钥。

[default]
aws_access_key_id = <Your Volces AK>
aws_secret_access_key = <Your Volces SK>

执行以下命令保存配置。

:wq

运行以下代码，使用 iterable-style datasets。

说明

本文以使用北京地域的 bucketname 桶的 shard-data/ 目录为例，建议您根据业务需求修改代码。

from s3torchconnector import S3IterableDataset, S3ClientConfig

if __name__ == '__main__':
dataset_uri = 's3://bucketname/shard-data/'
region = 'your region'
endpoint = 'your endpoint'
s3_config = S3ClientConfig(force_path_style=False)

print('测试 iterable-style datasets')
iterable_dataset = S3IterableDataset.from_prefix(dataset_uri,
region=region,
endpoint=endpoint,
s3client_config=s3_config)
# 遍历并读取Dataset中的对象。
for item in iterable_dataset:
print(item.bucket, item.key)
data = item.read()
print(len(data))

运行以下代码，使用 map-style datasets。

说明

本文以使用北京地域的 bucketname 桶的 shard-data/ 目录为例，建议您根据业务需求修改代码。

from s3torchconnector import S3IterableDataset, S3ClientConfig

if __name__ == '__main__':
dataset_uri = 's3://bucketname/shard-data/'
region = 'your region'
endpoint = 'your endpoint'
s3_config = S3ClientConfig(force_path_style=False)

print('测试 map-style datasets')
map_dataset = S3MapDataset.from_prefix(dataset_uri,
region=region,
endpoint=endpoint,
s3client_config=s3_config)
# 遍历并读取Dataset中的对象。
for item in map_dataset:
print(item.bucket, item.key)
data = item.read()
print(len(data))

# 根据索引读取Dataset中的对象。
item = map_dataset[0]
bucket = item.bucket
key = item.key
print(bucket, key)
data = item.read()
print(len(data))

运行以下代码，保存或加载 checkpoint。

说明

本文以使用北京地域的 bucketname 桶的 checkpoint/ 目录为例，建议您根据业务需求修改代码。

import torch
import torchvision
from s3torchconnector import S3Checkpoint, S3ClientConfig

if __name__ == '__main__':
checkpoint_uri = 's3://bucketname/checkpoint/'
region = 'your region'
endpoint = 'your endpoint'
s3_config = S3ClientConfig(force_path_style=False)

print('测试 checkpointing')
checkpoint = S3Checkpoint(region=region,
endpoint=endpoint,
s3client_config=s3_config)

model = torchvision.models.resnet18()

with checkpoint.writer(checkpoint_uri + 'epoch0.ckpt') as writer:
torch.save(model.state_dict(), writer)

with checkpoint.reader(checkpoint_uri + 'epoch0.ckpt') as reader:
state_dict = torch.load(reader)

model.load_state_dict(state_dict)

对象存储