背景介绍
在进行机器学习的模型训练任务的时候,针对图像数据集的处理,常常会对数据集进行分块的操作;具体到模型结构中,在ViT框架中,把每一个图像看作是一个的patch,每一个patch可以当作是一个NLP领域的一个编码信息,从而把处理NLP任务的transformer架构迁移到CV领域,同时取得了非常不错的结果
那具体如何在代码中实现图片的网格化分块呢?einops库提供了非常方便的函数工具
einops库
Einops(EINstein OPerationS)是一个Python库,其目标是提供一个简单、清晰和强大的接口,用于处理和重塑多维数组。它旨在改善深度学习中Tensor操作的可读性和灵活性。Einops在PyTorch,TensorFlow和numpy中都可以使用。
Einops库的主要功能包括:
-
重塑(Reshape):
rearrange
函数可以根据指定的模式,重新安排或者合并原有的维度。例如,将维度 (C, H, W) 变为 (H, W, C)。 -
降维(Reduce):
reduce
函数可以对指定的维度求和,求平均等。 -
重复(Repeat):
repeat
函数可以将array或者Tensor在指定的维度上,进行复制操作。
任务设计
把下面这张卡比巴拉的图片分块成的形状
代码实现
# 依赖的库环境
import einops
import torch
import torchvision
import matplotlib.pyplot as plt
from torchvision.transforms import ToTensor
from PIL import Image
# 读取图片数据
img = Image.open(r"D:\deep_learning\1_27_ViT\kabibala.jpg")
# 对读取的PIL格式数据进行裁剪和格式转换
img = torchvision.transforms.Compose([torchvision.transforms.CenterCrop(1000),ToTensor()])(img)
# 对图片进行分割
img = einops.rearrange(img,'c (b1 h) (b2 w) -> (b1 b2) c h w',b1=10,b2=10)
# 图片网格化
img = torchvision.utils.make_grid(img,nrow=10,pad_value=0.9,padding=5)
# 图片显示
plt.imshow(img.permute(1,2,0))
plt.show()
代码中,img = einops.rearrange(img,'c (b1 h) (b2 w) -> (b1 b2) c h w',b1=10,b2=10)是对img数据的格式进行更改,原来图片的形状定义为[c,(b1 h),(b2 w)],指定b1=b2=10,就是在像素宽和像素高两个维度上设定了10的参数,'->'后的形状是变换之后的数据形状,将b1和b2转换到最前面的维度,得到[(b1 b2) c h w],即得到100张[c h w]的分割图片,这里可以等效于b1 b2 = b批次,所以可以使用torchvision.utils.make_grid将多批次的图像进行网格化的格式存储,最后的图片为
如果是模型训练的过程,则可以直接将[(b1 b2) c h w]格式的数据作为多批次的训练数据进行进一步的调整细化,以满足模型训练的要求
欢迎大家讨论交流~