Training on complete MMLV levels

[schrum2]

This might be overly ambitious, but I want to experiment with training on samples that represent entire MMLV levels. @almzayena-lab will provide the instructions for downloading a range of MMLV levels and converting to VGLC (the instructions might already be in a README somewhere she can point out).

What @fwilson12 will be doing is something similar to what @patwikk is doing here:
schrum2/MarioMakerPCG#5

First analyze the data to see what sort of range of sizes we have, and then we can discuss options to padding certain levels to make training data.

[fwilson12]

Just began captioning the dataset Alaa provided on slack even though we noticed some weird patterns/tile mappings in the data. Will also be tracking the runtime (~5000 samples from 300 levels) so we have a good idea of how long captioning sizeable datasets will be.

[fwilson12]

Also-- here's the level dimensionality scatterplot generated from the visualization script I ported from the MarioMakerPCG repo:

(from first 300 downloaded levels starting at id 200,000)

[schrum2]

Got it. If we want to train on complete levels, we need to bucket these into level gorups of comparable size, and pad all smaller levels within that bucket to a common size. We also want to fill the buckets to be of comparable size. I'm not sure of the best way to do that. Maybe just divide the total number of levels by 5 (or some number) and aim to have roughly that many levels per bucket. Then start expanding the width and heigt along a diagonal from 0,0 to the top right until we encompass that many levels, then keep going until you capture that many more, and so on until all are captured.

Also, I'm trying to think of ways to give us more flexibility, and have less onerous data wrangling. The data batches are handled by the level dataset, and by the modified bucket sampler you made. From what I understand, the assumption there is that all scenes in a bucket have the same size, but that means we have to make sure the data is sampled correctly.

What if each bucket sets a size, but we store (or at least input) the data at its original size, and then the code is responsible for padding levels to the shared bucket size? I think this will spare us a lot of hassle on the dataset creation end.

Also, in anticipation of some inventive options down the road, we should support two approaches to padding. The one that is easier to implement, but will probably lead to worse results, is to just pad the smaller levels with some empty space tile, or in the case of Mega Man, the null/void tile. Because it is easier, we'll start with that. However, the second approach is to implement loss masks and instead input a separate one-layer binary tensor that indicates whether or not a given tile position in the level is actually part of the level or not. Based on this input, there is a way to make the network ignore the masked locations during training. This also means that at inference, the shape of the output will match the shape of the non-masked area. For the moment, my intent is that this approach still allows null tiles, since that allows the diffusion model to decide how the level zig-zags within the available space.

Anyway, lot of information there. Whenever we eventually train on this, we'll probably want to output samples more frequently and focus on a small number of epochs.

[fwilson12]

Just began training an unconditional model on a dataset of ~200 MMLV levels with this command:

python train_diffusion.py python train_diffusion.py --complete_levels --json megaman\MMLV_Whole_LevelsAndCaptions-train.json --val_json megaman\MMLV_Whole_LevelsAndCaptions-validate.json --augment --game MM-Full --save_image_epochs 10 --output_dir megaman\MM-full-whole --tileset datasets\MM.json

Some of the levels are massive; this may only be feasible if we cap the maximum size of the levels we're training on; this first run should give us some good info to work with moving forward

Also: here's the new bucketing CLI output when training with a mixed-size dset:

Training samples: 181
remove_upside_down_pipes: False
bucket_levels: 181 levels grouped into 5 size bucket(s): ['80x28 (n=36)', '156x56 (n=36)', '216x72 (n=36)', '340x112 (n=36)', '800x280 (n=37)']
Loading data from megaman\MMLV_Whole_LevelsAndCaptions-validate.json...
Training samples: 9
remove_upside_down_pipes: False
bucket_levels: 9 levels grouped into 5 size bucket(s): ['16x16 (n=1)', '80x28 (n=2)', '80x44 (n=2)', '128x36 (n=2)', '144x56 (n=2)']

[fwilson12]

update: unsurprisingly, the run crashed w/ a CUDA error.

Looks we'll have to cap the max level area to some quantity such that the batch tensors aren't absurdly large.

I'll see if I can implement a cap in the bulk_downloader script that queries the mega man maker for levels, but I fear the majority of (good quality, coherent) levels will still be too large for the 4060s to process.

[fwilson12]

Here are the json files I attempted to train with:
https://drive.google.com/drive/folders/1Gz-kK9zvsTX_aJ-Qk9zTyjKYcDcPLkP1

[fwilson12]

updated command

python train_diffusion.py --complete_levels --json megaman\MMLV_Whole_LevelsAndCaptions-train.json --val_json megaman\MMLV_Whole_LevelsAndCaptions-validate.json --augment --game MM-Full --save_image_epochs 1 --output_dir megaman\MM-full-unconditional-whole --tileset datasets\MM.json --batch_size 5 --gradient_accumulation_steps 6

[schrum2]

I got the same memory error:

torch.AcceleratorError: CUDA error: out of memory
Search for `cudaErrorMemoryAllocation' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

I guess this means for us that training with complete levels simply is not possible, but I wonder what we can do that is as close to this as possible. Let's investigate a little more before giving up and closing the issue.

[fwilson12]

Surprisingly, the run on the lab machine was able to get to epoch 160 over the weekend; I ended up changing the batch size to 1 to be extra safe, so it seems the UNET was able to handle the largest level tensor on its own. What was unsurprising, though, is that the samples generated by the model were almost exclusively null tile:

The full MMLV levels we're pulling from MM Maker tend to have giant null segments, especially when they feature teleporters. If we still want to train on complete MMLV levels, we'll likely need to filter out levels with profound void discontinuities, or break them up into smaller level chunks (#198)

[schrum2]

Didn't we get better results than this at some point? Not great results ... but I thought that with a bit of rudimentary filtering we were able to train on full size levels and get something that was not filled with void.

[fwilson12]

These were the results from the first and only model I trained on full MMLV levels; maybe Alaa trained something on full levels?

We also trained a model on 64x64 VGLC data a few weeks ago, but the scenes contained some weird null segments due to the way we were creating datasets at the time, so maybe you're thinking of that?