基于存内计算架构的模型部署与映射优化

180 阅读6分钟

先进计算大赛背景

"存内计算"架构通过消除存储与计算单元间的物理距离,突破了传统冯·诺依曼架构的限制。自2016年起,该架构受到广泛关注,被视为国产算力发展的关键技术。

在存内计算架构中,权重布局对提高存算单元利用率和计算效率至关重要,是编译阶段优化的重点。优化算法的效率、阵列面积利用率和计算效率是神经网络在存内计算芯片上部署的关键指标。简而言之,"存内计算"技术通过优化权重布局,提升了芯片性能和资源利用效率。

存算一体技术要点

存算一体被认为是一种新型计算架构,与经典的冯诺依曼架构不同。在存算一体架构中,存储器本身即可进行计算,将存储单元和计算单元合为一体,省去了计算过程中数据搬运环节,消除了由于数据搬运带来的功耗和延迟,从而进一步提升计算能效。

image.png

以知存科技的模拟存算计算方案为例:

  1. 首先将被乘矩阵的参数提前存入存算单元
  2. 然后将乘数向量输入
  3. 在水平方向上,乘数在存算单元中完成与被乘数的乘法操作
  4. 在垂直方向上,每个存算单元的乘法结果累加
  5. 最后得到输出向量

之所以以矩阵计算为例,是因为矩阵对于机器学习系统而言极为关键,它为数据表示提供了一种既简单又高效的方法。举例来说,输入数据(如图像里的像素集合)或者模型内部不同层之间的运行机制均能够借助矩阵来表示。所以,矩阵相乘运算在深度学习模型的总计算量当中占据着相当大的比重。实际上,在诸多当下流行的Transformer模型,如BERT、CLIP以及ChatGPT中,矩阵乘法的运行时长大约占其总运行时长的45 - 60%。矩阵乘法在卷积运算的计算过程里起着重要的作用,而卷积运算又是大多数计算机视觉模型的基础,也是许多高性能计算应用的核心部分。

使用的create函数

python

def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
    from pathlib import Path
    from models.common import AutoShape, DetectMultiBackend
    from models.experimental import attempt_load
    from models.yolo import ClassificationModel, DetectionModel, SegmentationModel
    from utils.downloads import attempt_download
    from utils.general import LOGGER, ROOT, check_requirements, intersect_dicts, logging
    from utils.torch_utils import select_device

    # 函数实现细节省略...

    return model.to(device)

image.png

基于存算一体的深度学习编译工具链

深度学习编译器与传统编译器类似,是一种用于将深度学习神经网络模型部署到硬件平台的工具。它处在深度学习框架与硬件设备之间,把深度学习框架中所描述的模型定义当作输入内容,然后在各类深度学习硬件上生成高效的代码实现并将其作为输出结果。它是连接软硬件的桥梁,有着重要意义。

witin_mapper是知存科技自研的用于神经网络映射的编译软件栈,可以将量化后的神经网络模型映射到WTM2101 MPU加速器上。它是一种包括RiscV和MPU的完整解决方案,可以完成算子和图级别的转换和优化,将预训练权重编排到存算阵列中,并针对网络结构和算子给出存算优化方案。同时,它将不适合MPU运算的算子调度到CPU上运算,实现整网的调度,让神经网络开发人员高效快捷地将训练好的算法运行在WTM2101芯片上,极大缩短模型移植的开发周期并提高算法开发的效率。

更多信息可参考:Witmem-Toolchain-WTM2101

可用工具和工作流程

  1. 模型准备:使用官方提供的模型和通过官网下载的模型可视化工具Netron软件。
  2. 获取模型中数据依赖关系以及权重矩阵块参数。
  3. 对应onnx文件的权重矩阵块组信息以及矩阵块之间的数据处理先后关系。
  4. 进行编译映射工作,实现紧密排布。
  5. 输出结果,显示存算阵列情况与各个阵列内的排布情况。

输出结果包含以下信息:

  • core_id:"空矩形"的序号
  • index:环节一中解析得到的权重矩阵的数据依赖顺序
  • w_clo_start:权重矩阵的左上角在"空矩形"中的列方向位置
  • w_row_start:权重矩阵的左上角在"空矩形"中的行方向位置

配置验证

主要测试存算阵列的执行效率和排布的合理性,包括:

  • 是否超出存算阵列的范围
  • 是否存在矩阵块面积重叠的情况
  • 存算阵列的面积利用率

run函数

python

@smart_inference_mode()
def run(
    weights=ROOT / "yolov5s-cls.pt",
    source=ROOT / "data/images",
    data=ROOT / "data/coco128.yaml",
    imgsz=(224, 224),
    device="",
    view_img=False,
    save_txt=False,
    nosave=False,
    augment=False,
    visualize=False,
    update=False,
    project=ROOT / "runs/predict-cls",
    name="exp",
    exist_ok=False,
    half=False,
    dnn=False,
    vid_stride=1,
):
    # 函数实现细节省略...
  1. 函数参数:

    • weights: 模型权重文件路径
    • source: 输入源(可以是图片、视频或摄像头)
    • data: 数据集配置文件
    • imgsz: 推理图像大小
    • device: 运行设备(CPU 或 GPU)
    • 其他参数控制可视化、保存结果等选项
  2. 初始化:

python

source = str(source)
save_img = not nosave and not source.endswith(".txt")
is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
is_url = source.lower().startswith(("rtsp://", "rtmp://", "http://", "https://"))
webcam = source.isnumeric() or source.endswith(".streams") or (is_url and not is_file)
screenshot = source.lower().startswith("screen")

这部分代码确定输入源的类型(文件、URL、网络摄像头等)。

  1. 加载模型:

python

device = select_device(device)
model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
stride, names, pt = model.stride, model.names, model.pt
imgsz = check_img_size(imgsz, s=stride)

加载预训练模型并设置相关参数。

  1. 创建数据加载器:

python

if webcam:
    dataset = LoadStreams(source, img_size=imgsz, transforms=classify_transforms(imgsz[0]), vid_stride=vid_stride)
elif screenshot:
    dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt)
else:
    dataset = LoadImages(source, img_size=imgsz, transforms=classify_transforms(imgsz[0]), vid_stride=vid_stride)

根据输入源类型创建适当的数据加载器。

  1. 推理循环:

python

for path, im, im0s, vid_cap, s in dataset:
    # 预处理
    im = torch.Tensor(im).to(model.device)
    im = im.half() if model.fp16 else im.float()
    
    # 推理
    results = model(im)
    
    # 后处理
    pred = F.softmax(results, dim=1)
    
    # 处理预测结果
    for i, prob in enumerate(pred):
        # 获取 top5 预测结果
        top5i = prob.argsort(0, descending=True)[:5].tolist()
        
        # 保存或显示结果
        if save_img or view_img:
            # 在图像上添加文本注释
        if save_txt:
            # 将结果保存到文本文件

这是函数的核心部分,它遍历数据集,对每个输入进行预处理、推理和后处理。

  1. 结果输出: 函数最后会输出处理速度信息,并保存结果(如果指定)。