How to retrain(fine-tuned) custom datasets

[Umal-exvc]

I am a Japanese university student studying Text-to-Image. I have read #307 and created a custom dataset, but I don't know what to do next. Can you give me some hints?

[borisdayma]

Hi,

To train the model you need to use the script tools/train/train.py.

You can use a command such as:

python train.py \
    --assert_TPU_available \
    --do_train --do_eval \
    --config_name config/mini_glu --dtype bfloat16 \
    --tokenizer_name boris/dalle-mini-tokenizer \
    --dataset_repo_or_path PATH_TO_DATASET \
    --blank_caption_prob 0.1 \
    --output_dir output --overwrite_output_dir \
    --per_device_train_batch_size 34 --gradient_accumulation_steps 3 \
    --optim distributed_shampoo --block_size 1024 --beta1 0.9 --beta2 0.99 \
    --learning_rate 0.0001 --warmup_steps 2000 \
    --mp_devices 1 \
    --logging_steps 100 --eval_steps 400 --save_steps 4000 --log_model

This command should work well on a TPU v3-8. You'd need to adapt it slightly if you're on a GPU.

If you use a trained model, replace config_name with model_name_or_path and add a reference to a local path, huggingface hub path, or wandb artifact (you can use the small model from the inference notebook for example).

I should probably add it in the README.md somewhere…

[kartikpodugu]

which framework is used for training ? I didn't find any torch or tf import in train.py ?