如何拆分棋盘图像为张量实现批量分析提升代码运行速度

逐格循环调用模型的写法没有利用深度学习框架的批量并行推理能力，每次单样本前向会产生大量框架调度、GPU kernel启动的额外开销，运行效率很低。只需要把所有预处理完的棋格张量拼接成一个批次，一次性输入模型就能大幅提升速度，且不会改变原有识别精度。

核心修改逻辑 #

• 保留原有的棋盘拆分、图像transform预处理逻辑不变
• 循环阶段仅完成单棋格的预处理操作，不执行单样本推理，将处理好的单棋格张量存入列表
• 调用torch.stack将列表内所有棋格张量拼接为带batch维度的总张量，形状为[64, 通道数, 棋格高度, 棋格宽度]
• 一次性将批次张量输入模型，拿到全部64个棋格的预测结果后，再逐个映射为对应棋子类别

优化后可直接运行的代码 #

import torch
from PIL import Image
from torchvision import transforms

squares = image_processing('/Users/Me/Downloads/Screen Shot 2022-07-04 at 11.09.04 AM.png')
image_transforms = transforms.Compose([  
                    transforms.ToTensor()])

square_tensor_list = []
for square in squares:
    square = Image.fromarray(square)
    square_tensor = image_transforms(square).float()
    square_tensor_list.append(square_tensor)

# 拼接为批次张量执行批量推理
batch_tensor = torch.stack(square_tensor_list, dim=0)
# 若使用GPU推理，取消下一行注释即可
# batch_tensor = batch_tensor.cuda()
output = model(batch_tensor)
_, predicted = torch.max(output.data, 1)

# 映射类别得到最终棋盘状态
board = [classes[pred.item()] for pred in predicted]

性能参考 #

• GPU环境下，该写法速度通常是原逐格推理的6~15倍
• CPU环境下也能减少大量框架重复调度的开销，获得2~4倍的速度提升
• 如果后续需要批量处理多张棋盘图片，可以将多张图的棋格张量凑成更大的批次，推理效率还能进一步提升

内容的提问来源于stack exchange，提问作者Ozzy08