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

前提条件

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

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

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

已安装 Python 3.8 及以上版本。

环境准备

对接 TOS 前，请先在 Python 3.8+ 环境安装 Pytorch，详情请参见 Pytorch。

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

对接方式说明

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

操作步骤

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

使用 fsspec 对接

使用 FSX 对接

使用 S3 Connector 对接

使用 tosutil 对接

执行以下命令安装 tosfs。

# 安装 tosfs for fsspec
pip install tosfs

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

from tosfs.core import TosFileSystem

# 配置访问密钥等信息
if __name__ == '__main__':
tosfs = TosFileSystem(
key='<Your Volces AK>',
secret='<Your Volces SK>',
endpoint_url='http://tos-cn-beijing.volces.com',
region='cn-beijing',
)

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

说明

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

from torchdata.datapipes.iter import FSSpecFileLister, FSSpecFileOpener

if __name__ == '__main__':
kwargs = {
'key' : '<Your Volces AK>',
'secret' : '<Your Volces SK>',
'endpoint_url' : 'https://tos-cn-beijing.volces.com',
'region' :'cn-beijing'
}

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
from tosfs.core import TosFileSystem

if __name__ == '__main__':
tosfs = TosFileSystem(
key='<Your Volces AK>',
secret='<Your Volces SK>',
endpoint_url='http://tos-cn-beijing.volces.com',
region='cn-beijing',
)

bucket = 'bucketname'
print('测试 checkpointing')
model = torchvision.models.resnet18()

# 将 tos 封装为 python 的标准读写接口
with tosfs.open(bucket+ '/checkpoint/epoch1.ckpt', 'wb') as writer:
torch.save(model.state_dict(), writer)

with tosfs.open(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 = 'cn-beijing'
endpoint = 'http://tos-s3-cn-beijing.volces.com'
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 = 'cn-beijing'
endpoint = 'http://tos-s3-cn-beijing.volces.com'
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 = 'cn-beijing'
endpoint = 'http://tos-s3-cn-beijing.volces.com'
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)

对象存储