lucidrains / magvit2-pytorch

Implementation of MagViT2 Tokenizer in Pytorch
MIT License
501 stars 30 forks source link

Large scale training #17

Open iejMac opened 8 months ago

iejMac commented 8 months ago

Hey, wanted to start this comm channel as I'm looking to do a large scale training run using some of this code. I'm happy to share graphs/samples as I go along and wanted to ask a few things to start off:

As always, thanks for this! Very helpful

lucidrains commented 8 months ago

@iejMac oh hey! yea, should be complete, probably one or two more bugs left to iron out

i would stick with LFQ initially, as that was what the magvit2 paper proposed, although some people have reported better results with FSQ. i put it into one repository so we can test them against each other and find out

jpfeil commented 8 months ago

Hi @iejMac I'd like to follow along if that's okay. It would be great if you could share any changes you make to the codebase to allow for larger scale training. I'm happy to share any weights I generate to help people get started with pretrained models.

iejMac commented 8 months ago

@jpfeil will do!

Ok I think I'm mostly set up (had to port this code to a repo with a different style). My first question is - do we have some prepared configs (like what layers, how many frames, what fps etc.) which roughly correspond to some models they trained in the paper? Just so we can compare.

For reference, currently I'm using the equivalent of this:

tokenizer = VideoTokenizer(
    image_size = 128,
    init_dim = 64,
    max_dim = 512,
    layers = (
        'residual',
        'compress_space',
        ('consecutive_residual', 2),
        'compress_space',
        ('consecutive_residual', 2),
        'linear_attend_space',
        'compress_space',
        ('consecutive_residual', 2),
        'attend_space',
        'compress_time',
        ('consecutive_residual', 2),
        'compress_time',
        ('consecutive_residual', 2),
        'attend_time',
    )
)
iejMac commented 8 months ago

Oh I also noticed one thing - is there a reason we don't normalize the pixels before passing it into the model? Or did I just not catch where that's done?

lucidrains commented 8 months ago

@iejMac oh hey, what is the typical normalization for video? i think .ToTensor() here should bring it to [0, 1]?

lucidrains commented 8 months ago

@iejMac are you using the LFQ from this repo? the main claim of this paper is that this new quantization method helps them scale to more codes and better generation scores. if i had to sum up the paper, it would be, use independent binary latents + mostly convolutions

iejMac commented 8 months ago

@lucidrains ah yeah ToTensor does but your VideoDataset doesn't do that and thats what I was using to test (was getting loss ~O(1e5)). https://github.com/lucidrains/magvit2-pytorch/blob/b2f105b9a398264c523c2088f20b7b468200ab78/magvit2_pytorch/data.py#L159

iejMac commented 8 months ago

Yes I'm using LFQ from that. The main question I have about config is like can we figure out a parametrization of VideoTokenizer (given all params you added) that corresponds to like MAGVIT2-small so we can do some nice test runs.

Let's start out with - 8 frame videos at 25 FPS. Given that what are reasonable params for layers and other values in order to get decent results.

With the setup I sent above the loss curve/reconstructions look like this and it usually gets a 'nan' at some point (that's where it ended):

Screenshot 2023-11-26 at 3 45 41 PM Screenshot 2023-11-26 at 3 45 59 PM
lucidrains commented 8 months ago

@iejMac shoot, i normalized for gifs, but not mp4s.. thank you Maciej!

lucidrains commented 8 months ago

@iejMac yup, i can get some of the hyperparameters inline with the paper's probably Tuesday (currently in the middle of another project)

iejMac commented 8 months ago

cool, was just wondering if you have something on hand. I'll try to read/play around and I'll report here if I come up with something. Also for lowish-effort video dataloading video2numpy could be a good option! It's pretty fast and does all the normal preprocessing for you. Maybe I'll make a PR for that if you're interested.

lucidrains commented 8 months ago

@iejMac would greatly appreciate it! 🙏

mudtriangle commented 8 months ago

Following in this thread since it's also related to video loading: Shouldn't there be a frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) inside the video_to_tensor function?

lucidrains commented 8 months ago

@mudtriangle there's a BGR format? 😅

mudtriangle commented 8 months ago

Yep, and I think it's still default in cv2.

lucidrains commented 8 months ago

@mudtriangle got it, put in a quick fix here (just doing it in tensor space, as i'm not familiar with cv2 enough)

Jason3900 commented 4 months ago

@mudtriangle got it, put in a quick fix here (just doing it in tensor space, as i'm not familiar with cv2 enough)

Hello, I'm wondering if there's any progress of hyperparameter/architecture config alignment with the magvit-v2 paper.