我想实现Vision Transformer模型，在论文中，他们表示他们将输入图像分割成一定分辨率的小块，例如，如果图像64x64和补丁分辨率为16x16，它将被分割成16个小补丁，每个小补丁的分辨率为16x16，因此最终形状为（N，P，P，C），其中N是贴片的数量，P是分辨率，C是通道的数量。
我尝试了什么，所以分裂是矢量化的：

def image_to_patches_fast(image, res_patch):
    
    (H, W, C) = get_image_shape(image)
    
    
    if C == 1:
        image = image.convert('RGB')
        (H, W, C) = get_image_shape(image)
                    
    P = res_patch
    N = (H*W)//(P**2)
        
    image_tensor = torchvision.transforms.PILToTensor()(image).permute(1,2,0)
    image_patches = image_tensor.view(N,P,P,C)

函数工作，但输出不是预期的，因为当我试图可视化补丁时，有一些错误，补丁可能没有很好地定位，或者我不知道，这里有一个例子：
输入图像：

输出面片的可视化：

可视化补丁的功能：

def show_patches(patches):
    
    N,P = patches.shape[0], patches.shape[1]
       
    nrows, ncols = int(N**0.5),int(N**0.5)
    fig, axes = plt.subplots(nrows = nrows, ncols=ncols)
    for row in range(nrows):

        for col in range(ncols):

            idx = col + (row*nrows)
            
            axes[row][col].imshow(patches[idx,:,:,:])
            axes[row][col].axis("off")

    plt.subplots_adjust(left=0.1,
                    bottom=0.1,
                    right=0.9,
                    top=0.9,
                    wspace=0.1,
                    hspace=0.1)
    plt.show()

我尝试了另一个函数来分割图像，但它的速度较慢，因为它使用循环，它的工作原理与预期一样：

def image_to_patches_slow(image, res_patch):
    
    (H, W, C) = get_image_shape(image)
    
    
    if C == 1:
        image = image.convert('RGB')
        (H, W, C) = get_image_shape(image)
                    
    P = res_patch
    N = (H*W)//(P**2)
    
    nrows, ncols = int(N**0.5), int(N**0.5)
    
    image_tensor = torchvision.transforms.PILToTensor()(image).permute(1,2,0)
    image_patches = torch.zeros((N,P,P,C),dtype = torch.int)
    
    
    for row in range(nrows):
        s_row = row * N
        e_row = (row * N) + N
        for col in range(ncols):

            idx = col + (row*nrows)

            s_col = col*N
            e_col = (col*N) + N
                
            image_patches[idx] = image_tensor[s_row:e_row, s_col:e_col]
    
    return image_patches

它的输出：

因此，任何帮助，因为这个缓慢的版本瓶颈的训练。