Video ViT Token Count

Problem

Count the tokens a video ViT consumes.

A video of t frames is cut into non-overlapping tubelets of tube_t frames, and each tubelet is patchified into (h // patch) * (w // patch) spatial tokens. Add one CLS token when include_cls is True.

Hint

Time is tokenized alongside space, so the tubelet grid is (t // tube_t) x (h // patch) x (w // patch).

Test cases

case 1

in: [16,224,224,2,16,true]

expected: 1569

case 2

in: [8,112,112,4,16,false]

expected: 98

case 3

in: [1,32,32,1,8,true]

expected: 17

case 4

in: [32,64,96,4,16,false]

expected: 192

Starter codePython

def video_token_count(t, h, w, tube_t, patch, include_cls):
    # Your code here
    pass

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