Before training you need to modify the *.YAML file path in the config folder, such as:
data_root: '../FSSD-12/'
train_list: './data_list/train/fold0_defect.txt'
val_list: './data_list/val/fold0_defect.txt'##VAE: You need download pretrained VAE From Huggingface.Set its weight path in train_kl_GP.py/train_kl_GP_tiny.py. such as:
vae = AutoencoderKL.from_pretrained("/PATH_TO_YOUR_VAE")
vae = AutoencoderTiny.from_pretrained("/PATH_TO_YOUR_VAE_TINY", torch_dtype=torch.float32).cuda()##Environment: This training framework is base on https://github.com/VDT-2048/CPANet. You can refer it to this project to set up your virtual environment.
##Running:
CUDA_VISIBLE_DEVICES=0 python train_kl_GP.py --config ./config/fold0_resnet50.yaml