magvit2-pytorch icon indicating copy to clipboard operation
magvit2-pytorch copied to clipboard

Large scale training

Open iejMac opened this issue 1 year ago • 17 comments

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:

  • Is the implementation of the original paper functionality complete?
  • What is the best configuration you have found? (I'm seeing talks about LFQ vs FSQ and I see code for diff transformers etc.)

As always, thanks for this! Very helpful

iejMac avatar Nov 22 '23 05:11 iejMac

@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

lucidrains avatar Nov 22 '23 15:11 lucidrains

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.

jpfeil avatar Nov 22 '23 18:11 jpfeil

@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 avatar Nov 26 '23 08:11 iejMac

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?

iejMac avatar Nov 26 '23 08:11 iejMac

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

lucidrains avatar Nov 26 '23 17:11 lucidrains

@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

lucidrains avatar Nov 26 '23 18:11 lucidrains

@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 avatar Nov 26 '23 23:11 iejMac

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

iejMac avatar Nov 26 '23 23:11 iejMac

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

lucidrains avatar Nov 26 '23 23:11 lucidrains

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

lucidrains avatar Nov 26 '23 23:11 lucidrains

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.

iejMac avatar Nov 26 '23 23:11 iejMac

@iejMac would greatly appreciate it! 🙏

lucidrains avatar Nov 27 '23 00:11 lucidrains

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?

mudtriangle avatar Nov 27 '23 17:11 mudtriangle

@mudtriangle there's a BGR format? 😅

lucidrains avatar Nov 27 '23 17:11 lucidrains

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

mudtriangle avatar Nov 27 '23 17:11 mudtriangle

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

lucidrains avatar Nov 27 '23 18:11 lucidrains

@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.

Jason3900 avatar Mar 07 '24 02:03 Jason3900