Training on a new dataset

[google-deepmind]

Training on a new dataset

[cdoersch]

Kubric uses the same format as the other datasets: a set of query_points in t,y,x format (shape [batch, num_points, 3]), a set of target_points in x,y format (shape [batch, num_points, num_frames, 2]), and an binary occlusion flag (1 if occluded, 0 otherwise; shape [batch, num_points, num_frames]), and the video (scaled between -1 and 1 shape [batch, num_frames, height, width, 3]). To train on something else, you'd need to edit experiment.py to add a new dataset_constructor which will return a new python generator that generates dicts containing the above fields. Then you need to edit the config file to add the desired dataset name and its kwargs to datasets. That is, once you've written the generator, it should only be a few lines of code.

Note that the code is set up for multi-dataset training, meaning that the training class will receive a dict keyed by dataset name, with an example from each dataset. You may need to change the input_key in supervised_point_prediction.py to get it to use the correct dataset.

[chandlj]

Thanks for the reply, that worked well for me. What config parameters did you set for finetuning, such that I could finetune the existing model on the new dataset I added? Right now I'm using the checkpoint.npy file that is available in the README, but that already initializes to having run 100k steps. In my config, I have to set the number of training_steps to be something >100k to get it to run. In the paper you mention that for finetuning you ran 5000 steps with 100 warmup steps and a learning rate of 1e-5, how do I set those parameters to be compatible with the existing checkpoint.npy file? Did you change the weight decay parameters at all?

[aloma85]

Can I use my personal dataset that is not related to any of the ones mentioned above?