train.py

import torch
import argparse
import os
import numpy as np
import torch.nn.functional as F
from datasets.mvtec_supervised import MVTecDataset
from datasets.visa_supervised import VisaDataset
import models.vv_open_clip as open_clip
import torchvision.transforms as transforms
from utils.loss import FocalLoss, BinaryDiceLoss
from models.FiLo import FiLo
from tqdm import tqdm
from prefetch_generator import BackgroundGenerator


class DataLoaderX(torch.utils.data.DataLoader):
    def __iter__(self):
        return BackgroundGenerator(super().__iter__())

mvtec_obj_list = [
    "bottle",
    "cable",
    "capsule",
    "carpet",
    "grid",
    "hazelnut",
    "leather",
    "metal nut",
    "pill",
    "screw",
    "tile",
    "toothbrush",
    "transistor",
    "wood",
    "zipper",
]

visa_obj_list = [
    "candle",
    "cashew",
    "chewinggum",
    "fryum",
    "pipe fryum",
    "macaroni1",
    "macaroni2",
    "pcb1",
    "pcb2",
    "pcb3",
    "pcb4",
    "capsules",
]

positions_list = ['top left', 'top', 'top right', 'left', 'center', 'right', 'bottom left', 'bottom', 'bottom right']


if __name__ == "__main__":

    parser = argparse.ArgumentParser("FiLo Train", add_help=True)
    parser.add_argument(
        "--clip_model", type=str, default="ViT-L-14-336", help="model used"
    )
    parser.add_argument(
        "--clip_pretrained",
        type=str,
        default="openai",
        help="pretrained weight used",
    )
    parser.add_argument(
        "--features_list",
        type=int,
        nargs="+",
        default=[6, 12, 18, 24],
        help="features used",
    )

    parser.add_argument(
        "--train_data_path",
        type=str,
        default="./data/visa",
        help="train dataset path",
    )
    parser.add_argument("--image_size", type=int, default=518, help="image size")

    parser.add_argument(
        "--dataset", type=str, default="visa", help="train dataset name"
    )
    parser.add_argument("--aug_rate", type=float, default=0.2, help="image size")
    parser.add_argument("--batch_size", type=int, default=1, help="batch size")

    parser.add_argument(
        "--learning_rate", type=float, default=0.001, help="learning rate"
    )
    parser.add_argument(
        "--decoder_learning_rate", type=float, default=0.0001, help="learning rate for decoder"
    )
    parser.add_argument(
        "--adapter_learning_rate", type=float, default=0.00001, help="learning rate for adapter"
    )
    parser.add_argument("--epoch", type=int, default=15, help="epochs")
    parser.add_argument("--n_ctx", type=int, default=12, help="epochs")
    parser.add_argument("--adapter_epoch", type=int, default=5, help="epochs")

    parser.add_argument(
        "--save_path",
        type=str,
        default="./ckpt",
        help="path to save results",
    )

    parser.add_argument(
        "--device", type=str, default="cuda", help="running on cpu only!, default=False"
    )
    args = parser.parse_args()


    save_path = args.save_path
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    device = args.device
    image_size = args.image_size
    batch_size = args.batch_size
    epochs = args.epoch
    dataset_name = args.dataset

    transform = transforms.Compose(
        [
            transforms.Resize((image_size, image_size)),
            transforms.CenterCrop(image_size),
            transforms.ToTensor(),
        ]
    )

    _, _, preprocess = open_clip.create_model_and_transforms(
        args.clip_model, image_size, pretrained=args.clip_pretrained
    )

    if dataset_name == "visa":
        train_data = VisaDataset(
            root=args.train_data_path,
            transform=preprocess,
            target_transform=transform,
        )
    else:
        train_data = MVTecDataset(
            root=args.train_data_path,
            transform=preprocess,
            target_transform=transform,
            aug_rate=args.aug_rate,
        )

    train_dataloader = DataLoaderX(
        train_data, batch_size=batch_size, shuffle=True, num_workers=8
    )
    
    obj_list = [x.replace("_", " ") for x in train_data.get_cls_names()]

    filo_model = FiLo(obj_list, args, device).to(device)
    
    main_part_param_groups = [
        {'params': filo_model.decoder_cov.parameters(), 'lr': args.decoder_learning_rate},
        {'params': filo_model.decoder_linear.parameters(), 'lr': args.decoder_learning_rate},
        {'params': filo_model.normal_prompt_learner.parameters(), 'lr': args.learning_rate},
        {'params': filo_model.abnormal_prompt_learner.parameters(), 'lr': args.learning_rate}
    ]

    optimizer_main_part = torch.optim.AdamW(
        main_part_param_groups,
        betas=(0.5, 0.999),
    )

    adapter_param_groups = [
        {'params': filo_model.adapter.parameters(), 'lr': args.adapter_learning_rate},
    ]

    optimizer_adapter = torch.optim.AdamW(
        adapter_param_groups,
        betas=(0.5, 0.999),
    )

    # losses
    loss_focal = FocalLoss()
    loss_dice = BinaryDiceLoss()

    with torch.no_grad():
        obj_list = [x.replace("_", " ") for x in train_data.get_cls_names()]


    for epoch in range(epochs):
        loss_list = []
        for items in tqdm(train_dataloader):
            image = items["img"].to(device)
            cls_name = items["cls_name"][0]
            image_path = items["img_path"]
            anomaly_cls = items["anomaly_class"][0]
            label = items['anomaly'].to(device)
            text_probs, anomaly_maps = filo_model(items, with_adapter=False)

            # losses
            gt = items["img_mask"].squeeze().to(device)
            gt[gt > 0.5], gt[gt <= 0.5] = 1, 0
            loss = 0
            for num in range(len(anomaly_maps)):
                loss += loss_focal(anomaly_maps[num], gt)
                loss += loss_dice(anomaly_maps[num][:, 1, :, :], gt)
                loss += loss_dice(anomaly_maps[num][:, 0, :, :], 1 - gt)

            optimizer_main_part.zero_grad()
            loss.backward()
            optimizer_main_part.step()

            loss_list.append(loss.item())
        # logs
        if (epoch + 1) % 1 == 0:
            print(
                "epoch [{}/{}], loss:{:.4f}".format(
                    epoch + 1, epochs, np.mean(loss_list)
                )
            )

    for epoch in range(args.adapter_epoch):
        loss_list = []
        for items in tqdm(train_dataloader):
            image = items["img"].to(device)
            cls_name = items["cls_name"][0]
            image_path = items["img_path"]
            anomaly_cls = items["anomaly_class"][0]
            label = items['anomaly'][0].to(device)
            text_probs, anomaly_maps = filo_model(items, only_train_adapter=True, with_adapter=True)

            # losses
            text_probs = text_probs[:, 0, ...] / 0.07
            loss = F.cross_entropy(text_probs.squeeze(), label)
            loss_list.append(loss.item())

            optimizer_adapter.zero_grad()
            loss.backward()
            optimizer_adapter.step()

            loss_list.append(loss.item())

        # logs
        print(
            "adapter epoch [{}/{}], loss:{:.4f}".format(
                epoch + 1, args.adapter_epoch, np.mean(loss_list)
            )
        )

        # save mode
        save_name =  + "filo_train_on_" + args.dataset
        ckp_path = os.path.join(
            save_path,
            f"{save_name}.pth",
        )
        torch.save(
            {
                "filo": filo_model.state_dict(),
            },
            ckp_path,
        )