ViT Patch Embedding

Problem

Turn a square image into a sequence of flattened non-overlapping patches in row-major order, scanning patch rows top to bottom and patches left to right.

The signature is patch_embedding(image, patch_size). Each patch is flattened row-major into a vector of length patch_size^2.

Hint

Patch embedding is just a reshape plus flatten before the linear projection in ViT.

Test cases

case 1

in: [[[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14,15,16]],2]

expected: [[1,2,5,6],[3,4,7,8],[9,10,13,14],[11,12,15,16]]

case 2

in: [[[1,2],[3,4]],1]

expected: [[1],[2],[3],[4]]

case 3

in: [[[1,2,3],[4,5,6],[7,8,9]],3]

expected: [[1,2,3,4,5,6,7,8,9]]

case 4

in: [[[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14,15,16]],4]

expected: [[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16]]

Starter codePython

def patch_embedding(image, patch_size):
    # Your code here
    pass

4 test cases · run your code against them in the editor