用于MAX78000模型训练的数据加载器设计

MAX78000是具有超低功耗卷积神经网络加速器的人工智能微控制器，可以在芯片上有效地运行人工智能模型。用户应首先使用ADI公司在PyTorch上的开发流程开发神经网络模型。然后，MAX78000频率合成器工具接受YAML格式的PyTorch检查点和模型描述，自动生成C代码，在MAX78000上编译和执行。模型开发阶段使用的基本软件组件之一是数据加载器，它负责特定于应用程序的数据准备任务。本文档介绍在准备适合MAX78000模型训练的特定应用训练和验证/测试集实体时，数据加载器实现的原则和设计注意事项。

介绍

在应用程序开发周期中，第一步是准备和预处理可用数据以创建训练和验证/测试数据集。除了通常的数据预处理外，在MAX78000上运行模型还需要考虑几个硬件限制。

数据加载器的主要职责可以总结如下：

[可选]将原始资源的输入和标签数据下载到通过调用ADI公司的CNN培训工具（培训存储库/train.py提供的数据路径中。

从指定的数据路径（csv/二进制文件、带或不带层次结构的文件夹/s 等）读取原始输入数据。

读取提供的数据路径中的原始标签/注释（csv/二进制文件/s、带或不带层次结构的文件夹等）。

[可选]应用数据预处理步骤，如增强、数据清理等。

对输入数据和标签应用所需的数据类型和范围转换。

执行训练和测试/验证拆分。

提供数据加载器方法和与MAX78000模型训练工具兼容的定义字典。

[可选]将处理后的数据实体保留在磁盘上，以便将来访问。

[可选]将上述步骤应用于可从同一原始数据源生成的每个不同数据集变体。

提供两个用于训练和测试数据的 PyTorch 数据集。

以下各节提供了有关创建高效数据加载器的说明，以满足所需的功能并方便地集成训练工具。

图 1 抽象地显示了数据加载器实现的主流。以下各节介绍了详细信息。

图1.数据加载器模块的主流。

自定义数据加载器实现的设计原则

数据加载器实现的主要职责之一是在将数据集实体馈送到 CNN 模型之前进行数据范围调整和数据类型管理。图 2 总结了这些操作，以下各节将详细介绍这些操作。

图2.数据范围规范化和类型转换。

预期数据范围

对于训练，输入数据应在

.当评估量化权重或在硬件上运行时，输入数据应位于本机MAX7800X范围[-128， +127]。

如以下部分所述，数据加载器函数将数据路径和一些参数作为输入参数。参数字段包括两个必填字段：act_mode_8bit 和 truncate_testset。当设置为True时，第一个参数是指对于本地MAX7800X范围，即范围[-128， +127]，应正确进行大小写归一化。设置为 False 时，规范化应在训练范围内

如果可用数据在 [0 1] 范围内，例如，在 PIL 图像中，数据加载器可以直接调用 ai8x.normalize（） 函数，使用提供的 args 参数将数据规范化为两个支持的数据范围：

class normalize:
"""
Normalize input to either [-128/128, +127/128] or [-128, +127]
"""
    def __init__(self, args):
        self.args = args

    def __call__(self, img):
        if self.args.act_mode_8bit:
            return img.sub(0.5).mul(256.).round().clamp(min=-128, max=127)
        return img.sub(0.5).mul(256.).round().clamp(min=-128, max=127).div(128.)

如果可用数据范围为 [0 255]，则需要在调用 ai256x.normalize（） 函数之前将其除以 0 以使其达到 [1 8] 范围。

注意：ai8x 模块的设备设置方法ai8x.set_device也接受相关参数模拟：True 表示训练案例 （act_mode_8bit = True），False 表示量化模型的评估或在也初始化 act_mode_8bit = False 的硬件上运行。此方法由具有适当参数管理的训练脚本使用，但如果在外部调用函数，则应正确设置模拟参数。

在MAX7800X硬件上运行推理时，必须考虑本地数据格式，并且在推理过程中应尽可能少地进行预处理。

数据类型

数据源可能具有不同范围内各种格式和值的原始数据文件。数据集类和数据加载器函数负责处理必要的转换。

数据加载器函数应返回数据类的训练和测试数据集元组。类型转换和转换通常在 __get_item__ 函数内处理，该函数应返回指定索引数据实体的数据元组和标签。数据项的类型应为：火炬。库达]。形状火炬的浮动张量。大小（数据集字典的相关条目“输入”字段）。

标签维度可能因问题类型或输入数据形状而异。每个标签类型都应强制转换为 np.long，以便在训练脚本中正确计算训练损失。

在完成所有数据增强和预处理任务并将数据范围规范化为 [0 1] 后，应使用该ai8x_normalize进行适当的进一步规范化，然后可以使用 torchvision.transforms.ToTensor 执行类型转换。

注意：拿到火炬。库达]。FloatTensor，numpy 数组必须事先转换为 float32。

Torchvision软件包包括各种预处理转换，例如可以根据应用程序需求使用的PIL图像的随机裁剪。数据类可以利用Torchvision包的复合转换按顺序应用多个转换，例如ToTensor转换和ai8x_normalize每当访问数据条目时。

数据实体的存储

通常，有两种方法可以存储数据集条目;整个数据集条目可以存储在内存中，也可以在使用 __getitem__ 方法访问时从磁盘读取。基本的决策因素是数据集的大小和每个实体的大小。当数据集太大而无法放入内存时（在初始化函数中处理预处理和增强任务后），所有数据集条目都可以保存到磁盘中，并在以后每次访问时从磁盘单独读取。虽然将数据条目保留在内存中可以加快数据访问速度，但内存限制可能会阻止在所有情况下使用基于内存的方法。

注意：即使采用基于内存的方法，也建议将预处理和增强的数据条目写入磁盘，因为它们只执行一次。然后，在每次生成数据类实例时，可以执行将所有数据批量读取到内存中。

表 1 总结了同一数据源的两个数据加载器实现选项的一些度量。从第一行可以看出，磁盘存储方法可以处理更多图像。内存预算是限制已处理图像数量的因素。这两种方法的数据集生成时间都很长，因为第一种方法还处理预处理、扩充等步骤，然后将所有可用数据写入磁盘。基于内存的方法在以后生成数据集需要更长的时间，因为对象创建需要将大文件从磁盘批量读取到内存中。而在第二种方法中，每个数据集项都是独立保存的，__getitem__方法创建数据集对象和实体检索都花费很少的时间。第一种方法的内存消耗很高，因为它将所有数据集实体保留在内存中。在磁盘使用情况方面，第一种方法通常使用所有数据条目的单个文件，第二种方法为每个数据条目使用单独的文件。这应该会导致大致相似的磁盘预算。

注意：在表 1 中，由于处理图像的数量减少，第一种方法的磁盘空间要小得多。磁盘方法的唯一缺点是它增加了训练时间，因为每个数据输入读取都是作为单独的磁盘操作完成的。

	数据加载器 图像存储在内存中 图像存储在内存中	数据加载器 使用从磁盘读取的映像
可处理的图像数量	20 000 * 1 = 20000	34 426 * 3 = 103 278
数据集生成时间 – 首次运行	30 分	60 分
数据集生成时间 – 后续运行	15 分	瞬间
运行时内存消耗峰值	~55 千兆字节	~5 千兆字节
磁盘消耗	~50 千兆字节	~ 240 GB
训练时间 单纪元	60-62秒	1450 秒

自定义数据加载器实现的编程原则

数据加载器模块将在 PyTorch 中实现，预计至少具有以下三个组件：

数据集类定义。示例：类 AISegment（数据集）

torch.utils.data.Dataset是自定义数据集实现类应继承的抽象类。有关 PyTorch 中自定义数据加载器实现的教程，请参阅 [1]。

应重写__len__方法，以便 len（dataset） 返回数据集的大小。

还应该实现__getitem__来支持索引，以便数据集[i]可用于获取i千样本。对于MAX78000应用，该方法应返回一个数据元组及其相应的标签。

__init__功能参数和内容可根据应用需求进行定制。前两个参数通常是数据根路径和类型（测试或训练），如MAX78000训练存储库数据集文件夹中的几个数据加载器实现所示。但是，只要以下项目中介绍并作为外部通信点提供的数据加载器函数是固定的并执行所需的操作，就可以更改这些参数的顺序或命名。

数据加载器函数：不应修改此函数的签名。第一个输入是指定数据目录和程序参数的元组。两个前导布尔输入指定是否应加载训练和/或测试数据。

程序参数有两个与数据类实现相关的关键字段;act_mode_8bit和truncate_testset。第一个是指规范化类型（有关更多详细信息，请参阅预期数据范围部分），第二个用于将测试集截断为单个元素集。

示例：def AISegment352_get_datasets（data， load_train=True， load_test=True）。

数据集字典包括可用的数据加载器函数。字典名称和键值都不应更改，仅应根据自定义数据集实现调整值。同一数据源的每个变体都可以作为此字典中的单独元素存在。

例：

“name”键的值使ADI公司的CNN训练工具（训练存储库/train.py）能够在提供--dataset参数时查找数据集。因此，此字段的值在自定义数据集中应该是唯一的。

“input”键的值是输入数据的维度。第一个维度作为num_channels传递给模型，而其余维度作为维度传递。例如，“input”：（1， 28， 28） 以 num_channels=1 和 dimension=（28， 28） 的形式传递给模型。一维输入使用单个“维度”，例如“input”：（2， 512） 以 num_channels=2 和 dimension=（512， ） 的形式传递给模型。

“output”键的值指定分类问题的可用类类型。此键的值也可以使用字符串文本定义。

示例：“输出”：（“背景”、“纵向”）。

“weight”键的值指定每个数据实体的权重，并引用类标签。这是一个可选字段，如果未提供，则全部为“1”。

可以通过使用可用样本数为每个类提供成反比的权重来解决训练数据集中的类不平衡问题。因此，训练脚本更加关注频率较低的样本。

可选的回归可以设置为 True 以自动选择训练脚本的 --regression 命令行参数。

注意：当类数为 1 时，训练脚本会自动设置回归。 示例：“输出”：（“id”），“回归”：真。

示例数据加载器

MAX78000训练存储库数据集文件夹包括几种不同的数据加载器实现，详情请参见[2]。在本节中，将介绍一个定制的数据加载器来举例说明所有提到的原则。纵向分割数据集用于此目的，更多详细信息请参见 [3]。此数据集源包括 34，427 张分辨率为 600 × 800 的人类肖像图像（红色、绿色和蓝色 （RGB） 颜色格式），以及相同数量的标签图像，具有相同大小的红色、绿色、蓝色和 Alpha （RGBA） 格式的相应蒙版。

初始化

设计的数据加载器模块的第一个组件是具有以下初始化函数的数据加载器类。将跳过有关生成数据集信息数据框的详细信息。简而言之，这些行包括一些路径处理代码，用于保留原始图像路径、原始遮罩文件路径、裁剪 idx 和要保存到的数据集条目的泡菜文件路径。除了这些路径生成部分之外，初始化函数的主要功能是保留提供的参数并相应地排列一些局部变量（例如，训练或测试数据集信息数据帧）并生成数据集实体。

对于第一次初始化调用，所有数据处理任务都使用 __gen_datasets__ 方法处理，并为每个数据集项生成 pickle 文件并存储在磁盘上，以便在每次数据访问时读取。

class AISegment(Dataset):
…
    def __init__(self, root_dir, d_type, transform=None, im_size=[80, 80], fold_ratio=1):
        …
        self.d_type = d_type
        self.transform = transform
        self.img_ds_dim = im_size
        self.fold_ratio = fold_ratio

        # Generate and save dataset information file if not already available
        # Training and Test split is also performed here using the hash of file names (all three cropped images should fall into the same set)
        # Information data frames include raw data path, raw label path, crop idx, pickle file path, etc. for each data entity
        …

        # One of the created data frames is selected from: train_img_files_info & test_img_files_info
        if self.d_type == 'train':
            self.img_files_info = train_img_files_info

        elif self.d_type == 'test':
            self.img_files_info = test_img_files_info
        else:
            print('Unknown data type: %s' % self.d_type)
            return

   # Create and save pt files for each data entity (if not available before)
   self.__create_pt_files()
        self.is_truncated = False

    def __create_pt_files(self):
        if self.__check_pt_files_exist():
            return
        self.__makedir_exist_ok(self.processed_train_data_folder)
        self.__makedir_exist_ok(self.processed_test_data_folder)
        self.__gen_datasets()

数据增强

gen_datasets方法处理所有必需的预处理、扩充和预规范化步骤。实施的步骤如下：

从原始图像裁剪三个正方形图像（因为U-Net模型使用方形图像）。

根据提供的数据集参数，裁剪图像和遮罩图像被下采样为 80×80 或 352×352。

相应的遮罩图像将转换为二进制“背景”或“肖像”标签。

如果需要，将折叠图像（352×352 图像折叠为大小为 88×88×48 的图像）。

图像在保存之前按 256 缩放，因为复合变压器期望图像在 [0 1] 范围内，但原始图像是 RGB 并且值在 [0 255] 范围内。

图像裁剪原理如图3所示。

图 4 包括从同一原始图像构建的三个示例训练图像。

图3.示例数据加载器 – 数据增强：从原始图像中裁剪出三个方形图像。

图4.从原始图像（600 × 600）裁剪了三个大小为 800 × 600 的图像（以及相应的消光图像）。

gen_datasets方法的实现方式如下：

def __normalize_image(self, image):
        return image / 256

def __gen_datasets(self):

        # For each entry in dataset information dataframe
        for _, row in tqdm(self.img_files_info.iterrows()):

            img_file = row['img_file_path']
            matting_file = row['lbl_file_path']
            pickle_file = row['pickle_file_path']
            img_crp_idx = row['crp_idx']

            img = Image.open(img_file)
            lbl_rgba = Image.open(matting_file)

            vertical_crop_area = AISegment.img_dim[0] - AISegment.img_crp_dim[0]
            step_size = vertical_crop_area / (AISegment.num_of_cropped_imgs - 1)

            # Determine top left coordinate of the crop area
            top_left_x = 0
            top_left_y = 0 + img_crp_idx * step_size
                
            # Determine bottom right coordinate of the crop area
            bottom_right_x = AISegment.img_crp_dim[0]
            bottom_right_y = top_left_y + AISegment.img_crp_dim[0]
                
            img_crp = img.crop((top_left_x, top_left_y, bottom_right_x, bottom_right_y))
            img_crp_lbl = lbl_rgba.crop((top_left_x, top_left_y, bottom_right_x, bottom_right_y))
    
            img_crp = img_crp.resize(self.img_ds_dim)
            img_crp = np.asarray(img_crp).astype(np.uint8)

            img_crp_lbl = img_crp_lbl.resize(self.img_ds_dim)
            img_crp_lbl = (np.asarray(img_crp_lbl)[:, :, 3] == 0).astype(np.uint8) 
                
            # Fold the data (ex: 352 x 352 x 3 folded into 88 x 88 x 48) if required and save to pt file
            if self.fold_ratio == 1:
                img_crp_folded = img_crp

            else:
                img_crp_folded = None
                for i in range(self.fold_ratio):
                    for j in range(self.fold_ratio):
                        if img_crp_folded is not None:
                            img_crp_folded = np.concatenate((img_crp_folded, img_crp[i::self.fold_ratio, j::self.fold_ratio, :]), axis=2)
                        else:
                            img_crp_folded = img_crp[i::self.fold_ratio, j::self.fold_ratio, :]
            
            pickle.dump((img_crp_folded, img_crp_lbl), open(pickle_file, 'wb'))

数据加载器方法和转换器定义

数据加载器方法是第二个必需的组件定制数据加载器模块。对于示例 AISegment 数据集，实现了两个不同的数据加载器函数。第一个 （AISegment_get_datasets） 使用较小的 U-Net 网络模型返回大小为 80x80 的图像。后者 （AISegment352_get_datasets） 返回大小为 352×352 的图像。下面是第二个的实现，它生成具有所需属性的 AISegment 对象。复合变压器也在此函数中定义。此外，如果需要截断，则测试数据集将被截断。

def AISegment352_get_datasets(data, load_train=True, load_test=True):
    """…"""
    (data_dir, args) = data

    if load_train:
        train_transform = transforms.Compose([
            transforms.ToTensor(),
            ai8x.normalize(args=args)
        ])
        
        train_dataset = AISegment(root_dir=data_dir, d_type='train',
                                  transform=train_transform,
                                  im_size=[352, 352])
    else:
        train_dataset = None

    if load_test:
        test_transform = transforms.Compose([
            transforms.ToTensor(),
            ai8x.normalize(args=args)
        ])

        test_dataset = AISegment(root_dir=data_dir, d_type='test',
                                 transform=test_transform,
                                 im_size=[352, 352])

        if args.truncate_testset:
            test_dataset.data = test_dataset.data[:1]
    else:
        test_dataset = None

    return train_dataset, test_dataset

数据集字典

数据集字典是自定义数据加载器模块的第三个必需组件，该模块包含可用的数据加载器功能。对于示例 AISegment 数据加载程序，由于有两种数据集变体可以生成具有不同分辨率（80×80 或 352×352）的数据集实体，因此数据集字典有两个实体，每个实体都包括输入和输出大小的正确定义以及数据加载器函数名称。

使用图像测试训练模型

在模型开发阶段之后，可以使用测试数据集或任意样本来测试模型。关键的一点是，在向模型提供输入之前，必须在外部完成正确的转换/s数据加载器实现。

例如，使用AISegment数据集训练的示例模型需要输入形状[48， 88， 88]，这是分辨率为352×352的折叠RGB图像的通道优先表示，具有MAX7800X所需的归一化像素值。外部提供的测试图像甚至可能没有相同的颜色格式，但必须事先实现所需的转换，因为模型是针对 RGB 图像训练的。下面是具有 470×470 分辨率和 YbCr 颜色格式的纵向图像的测试模型的示例代码片段：

rgb_img = yuv_img.convert('RGB')
rgb_img_ds = rgb_img.resize([352, 352])

# Image to numpy array conversion:
rgb_img_ds = np.asarray(rgb_img_ds).astype(np.uint8)

# Fold image (352 x 352 x 3 folded into 88 x 88 x 48)
rgb_img_ds_folded = fold_image(rgb_img_ds, 4)

# Covert pixel values to range [0 1] and cast to float type (required for Torch)
rgb_img_ds_folded_scaled = (rgb_img_ds_folded / 256).astype(np.float32)

# Normalize for MAX78000
# Set act_mode_8bit=True as we will set model parameter simulate=True     
args = Args(act_mode_8bit=True)

transform = transforms.Compose([
            transforms.ToTensor(),
            ai8x.normalize(args=args)
])
rgb_img_ds_folded_scaled_normalized = transform(rgb_img_ds_folded_scaled)

# Add batch dimension
rgb_img_batch = rgb_img_ds_folded_scaled_normalized.unsqueeze(0)

# Load model
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
load_model_path = 'unet/qat_ai85unet_v7_352_4_best_q.pth.tar'

ai8x.set_device(device=85, simulate=True, round_avg=False)

model = mod.AI85Unet_v7_pt(num_classes=2, num_channels=3, dimensions=(88, 88), 
                        bias=True, fold_ratio=4)

checkpoint = torch.load(load_model_path, map_location=lambda storage, loc: storage)
ai8x.fuse_bn_layers(model)
model = apputils.load_lean_checkpoint(model, load_model_path, model_device=device)
ai8x.update_model(model)
model = model.to(device)

# Run model
with torch.no_grad():
    sample_img_rgb_batch = rgb_img_batch.to(device)
    model_out_rgb = model(sample_img_rgb_batch)

# Retrieve model output
out_vals = np.argmax(model_out_rgb[0, :, :, :].detach().cpu().numpy(), axis=0)
plt.imshow(out_vals, cmap='Greys')

图5包括以YCbCr格式给出的示例外部测试数据项，相应的RGB图像以及执行所有必需转换后的模型输出。首先，色彩空间需要转换为RGB。然后，应将图像缩减像素采样，使其具有 352 × 352 分辨率。下一个操作是折叠，需要转换和规范化。

图5.在 YCbCr 色彩空间、RGB 空间和模型输出中采样纵向图像。

审核编辑：郭婷