training on a different high dimensional discrete dataset

[adderbyte]

Hi, Thank you for the great work!

Any guidelines for training/testing on a different dataset -aside CIFAR and piano dataset?

Training on a different dataset would most likely involve chaning the data path in config file, any other adjustments to also keep in mind while also making these changes?

Thank you!

[andrew-cr]

Hi, thanks a lot for your interest in the work!

To run it on a new dataset you can make the following changes:

• Implement a new dataset class in lib/datasets/datasets.py, passing in parameters using the config object. The main points to note are having a cfg.data.S member that gives the number of possible discrete states and a cfg.data.shape that gives the shape of the data.
• For the model, if the data has structure similar to an image or a sequence you can use either the model used for CIFAR10 (UNet) or for the piano (transformer)
• For training, the code in lib/losses/losses.py will operate on tensors of shape B, D by flattening the input minibatch so if necessary you can add some code here if your dataset has a shape not currently implemented.
• For sampling you can use the classes in lib/sampling/sampling.py. This again operates on tensors of shape B, D so you just need to reshape the output back to the dataset shape after sampling finishes.
• If the problem is unconditional (like CIFAR10) you can use the loss GenericAux and if it is conditional (like the piano) you can use the loss ConditionalAux. The condition_dim parameter means that the model will be trained to predict the final D-condition_dim dimensions from the first condition_dim dimensions.

Please do let me know if you have any further questions on running it on the new dataset :)

[adderbyte]

Awesome! this is really helpful. Thank you for the detailed feedback.