train.py

import argparse

import datasets
import torch
import yaml
from torch.utils.data import DataLoader, IterableDataset
from torch.utils.tensorboard import SummaryWriter
from torchvision.transforms import v2 as transforms

import models
from dataset_misc import (
    ConvertedHuggingFaceIterableDataset,
    MVTecLOCOIterableDataset,
    TensorConvertedIterableDataset,
    TransformedIterableDataset,
)


def resnet_preprocess(image: torch.Tensor) -> torch.Tensor:
    preprocess = transforms.Compose(
        [
            transforms.Resize((512, 512), antialias=True),
        ]
    )
    return preprocess(image)


@torch.no_grad()
def find_distribution(
    model: torch.nn.Module,
    dataset: IterableDataset,
    device: torch.DeviceObjType,
    sample_size: int = 100,
) -> tuple:
    # find gaussian distribution of features
    preprocess = transforms.Compose(
        [
            transforms.Resize((512, 512), antialias=True),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ]
    )
    dataloader = DataLoader(TransformedIterableDataset(dataset, preprocess), batch_size=sample_size)

    batch = next(iter(dataloader)).to(device)
    features = model.forward(batch)
    features = torch.movedim(features, 0, 1)
    features = torch.flatten(features, start_dim=1)
    std, mean = torch.std_mean(features, dim=1, unbiased=False)

    return std, mean


@torch.no_grad()
def find_quantiles(
    teacher: torch.nn.Module,
    autoencoder: torch.nn.Module,
    student: torch.nn.Module,
    teacher_channels: int,
    dataset: IterableDataset,
    device: torch.DeviceObjType,
    sample_size: int = 1000,
) -> None:
    teacher.eval()
    autoencoder.eval()
    student.eval()

    preprocess = transforms.Compose(
        [
            transforms.Resize((256, 256), antialias=True),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ]
    )
    dataloader = DataLoader(TransformedIterableDataset(dataset, preprocess), batch_size=1)

    student_anomaly_maps = torch.empty((0, 64, 64), device=device)
    autoencoder_anomaly_maps = torch.empty((0, 64, 64), device=device)
    print("determining quantiles...")
    for image, i in zip(dataloader, range(sample_size)):
        image = image.to(device)
        teacher_result = teacher.forward(image)
        autoencoder_result = autoencoder.forward(image)
        student_result = student.forward(image)

        student_anomaly_map = torch.mean((teacher_result - student_result[:, :teacher_channels, :, :]) ** 2, dim=1)
        autoencoder_anomaly_map = torch.mean(
            (autoencoder_result - student_result[:, teacher_channels:, :, :]) ** 2,
            dim=1,
        )

        student_anomaly_maps = torch.cat((student_anomaly_maps, student_anomaly_map), dim=0)
        autoencoder_anomaly_maps = torch.cat((autoencoder_anomaly_maps, autoencoder_anomaly_map), dim=0)

        print(f"iteration: {i}/{sample_size}")

    quantiles = {
        "student_a": torch.quantile(student_anomaly_maps, q=0.9),
        "student_b": torch.quantile(student_anomaly_maps, q=0.995),
        "autoencoder_a": torch.quantile(autoencoder_anomaly_maps, q=0.9),
        "autoencoder_b": torch.quantile(autoencoder_anomaly_maps, q=0.995),
    }

    torch.save(quantiles, "models/quantiles.pt")

    print("finished determining quantiles!")


def train_teacher(
    channels: int,
    generic_dataset: IterableDataset,
    good_dataset: IterableDataset,
    device: torch.DeviceObjType,
    batches: int = 30_000,
    batch_size: int = 8,
    wideresnet_feature_layer: str = "layer1",
    wideresnet_feature_layer_index: int = 1,
    tensorboard_writer: SummaryWriter | None = None,
) -> torch.nn.Module:
    target_features_model = models.ChoppedWideResNet(
        channels=channels, layer_to_extract_from=wideresnet_feature_layer, layer_index=wideresnet_feature_layer_index
    ).to(device)
    target_features_model.eval()
    target_features_output_shape = target_features_model.forward(
        resnet_preprocess(torch.rand([1, 3, 1, 1]).to(device))
    ).shape

    print("finding normal distribution of features...")
    std, mean = find_distribution(target_features_model, generic_dataset, device, sample_size=50)
    std = std.unsqueeze(1).unsqueeze(1).expand(target_features_output_shape)
    mean = mean.unsqueeze(1).unsqueeze(1).expand(target_features_output_shape)

    common_preprocess = transforms.Compose(
        [
            transforms.Resize((512, 512), antialias=True),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            transforms.RandomGrayscale(0.1),
        ]
    )

    pdn_preprocess = transforms.Compose(
        [
            transforms.Resize(
                (target_features_output_shape[2] * 4, target_features_output_shape[3] * 4), antialias=True
            ),
        ]
    )

    dataloader = DataLoader(TransformedIterableDataset(generic_dataset, common_preprocess), batch_size=batch_size)

    teacher_pdn = models.PatchDescriptionNetwork(channels=channels).to(device)
    teacher_pdn.train()

    optimizer = torch.optim.Adam(teacher_pdn.parameters(), lr=1e-4, weight_decay=1e-5)

    print("starting training of teacher...")
    for image_batch, batch in zip(dataloader, range(batches)):
        image_batch = image_batch.to(device)

        with torch.no_grad():
            target_features = target_features_model.forward(resnet_preprocess(image_batch))
            target_features = torch.sub(target_features, mean)
            target_features = target_features / std
            target_features = torch.nan_to_num(target_features)

        predicted_features = teacher_pdn.forward(pdn_preprocess(image_batch))

        loss = torch.mean((target_features - predicted_features) ** 2)  # calculate mean square error
        print(f"batch: {batch}/{batches}  loss: {loss.item()}")
        if tensorboard_writer is not None:
            tensorboard_writer.add_scalar("teacher training", loss.item(), batch)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 1000 == 0:
            torch.save(teacher_pdn, "models/tmp/teacher.pt")

    torch.save(teacher_pdn, "models/generic_teacher.pt")
    print("finished training teacher!")

    print("normalizing teacher...")
    preprocess = transforms.Compose(
        [
            transforms.Resize((512, 512), antialias=True),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ]
    )
    dataloader = DataLoader(TransformedIterableDataset(good_dataset, preprocess), batch_size=batch_size)
    teacher_pdn = models.NormalizedPatchDescriptionNetwork(teacher_pdn).train().to(device)
    for image_batch, i in zip(dataloader, range(int(5_000 / batch_size))):
        teacher_pdn.forward(image_batch.to(device))
        print(f"batch: {i}/{int(5_000 / batch_size)}")

    torch.save(teacher_pdn, "models/teacher.pt")
    print("finished normalizing teacher!")

    return teacher_pdn


def train_autoencoder(
    channels: int,
    output_size: tuple[int, int],
    dataset: IterableDataset,
    device: torch.DeviceObjType,
    teacher: torch.nn.Module,
    epochs: int = 2_000,
    batch_size: int = 8,
    tensorboard_writer: SummaryWriter | None = None,
) -> torch.nn.Module:
    autoencoder = models.AutoEncoder(channels=channels, output_size=output_size).to(device)
    autoencoder.train()
    teacher.eval()

    optimizer = torch.optim.Adam(autoencoder.parameters(), lr=1e-4, weight_decay=1e-5)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=int(0.95 * epochs * len(dataset)), gamma=0.1)

    preprocess = transforms.Compose(
        [
            transforms.Resize((output_size[0] * 4, output_size[1] * 4), antialias=True),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            transforms.RandomGrayscale(0.1),
            transforms.RandomChoice(
                [
                    transforms.ColorJitter(brightness=0.2),
                    transforms.ColorJitter(contrast=0.2),
                    transforms.ColorJitter(saturation=0.2),
                ]
            ),
        ]
    )

    dataloader = DataLoader(TransformedIterableDataset(dataset, preprocess), batch_size=batch_size)

    print("starting training of autoencoder...")
    for epoch in range(epochs):
        for image_batch, batch in zip(dataloader, range(len(dataloader))):
            image_batch = image_batch.to(device)

            with torch.no_grad():
                teacher_result = teacher.forward(image_batch)
            autoencoder_result = autoencoder.forward(image_batch)

            loss = torch.mean((teacher_result - autoencoder_result) ** 2)

            total_batch = batch + epoch * len(dataloader)
            print(f"batch: {total_batch}/{epochs * len(dataloader)}  loss: {loss.item()}")
            if tensorboard_writer is not None:
                tensorboard_writer.add_scalar("autoencoder training", loss.item(), total_batch)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            scheduler.step()

        torch.save(autoencoder, "models/tmp/autoencoder.pt")

    torch.save(autoencoder, "models/autoencoder.pt")
    print("finished training autoencoder!")

    return autoencoder


def train_student(
    channels_teacher: int,
    channels_autoencoder: int,
    good_dataset: IterableDataset,
    generic_dataset: IterableDataset,
    device: torch.DeviceObjType,
    teacher: torch.nn.Module,
    autoencoder: torch.nn.Module,
    epochs: int = 3_000,
    batch_size: int = 8,
    tensorboard_writer: SummaryWriter | None = None,
) -> torch.nn.Module:
    student_pdn = models.PatchDescriptionNetwork(channels=channels_teacher + channels_autoencoder).to(device)
    student_pdn.train()
    teacher.eval()
    autoencoder.eval()

    optimizer = torch.optim.Adam(student_pdn.parameters(), lr=1e-4, weight_decay=1e-5)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=int(0.95 * epochs * len(good_dataset)), gamma=0.1)

    preprocess = transforms.Compose(
        [
            transforms.Resize(
                (256, 256), antialias=True
            ),  # quantile will not work with 512x512, TODO: sample random values for quantile calculation
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            transforms.RandomGrayscale(0.1),
        ]
    )

    dataloader_good = DataLoader(TransformedIterableDataset(good_dataset, preprocess), batch_size=batch_size)
    dataloader_generic = DataLoader(TransformedIterableDataset(generic_dataset, preprocess))

    print("starting training of student...")
    for epoch in range(epochs):
        for image_batch, generic_image, batch in zip(dataloader_good, dataloader_generic, range(len(dataloader_good))):
            image_batch = image_batch.to(device)
            generic_image = generic_image.to(device)

            with torch.no_grad():
                teacher_result = teacher.forward(image_batch)
                autoencoder_result = autoencoder.forward(image_batch)
            student_result = student_pdn.forward(image_batch)

            pdn_student_distance = (teacher_result - student_result[:, :channels_teacher, :, :]) ** 2
            pdn_student_quantile = torch.quantile(pdn_student_distance, q=0.999)
            pdn_student_hard_loss = torch.mean(pdn_student_distance[pdn_student_distance >= pdn_student_quantile])

            pdn_student_penalty_result = student_pdn.forward(generic_image)[:, :channels_teacher, :, :]
            pdn_student_penalty_loss = torch.mean(pdn_student_penalty_result**2)

            pdn_student_loss = pdn_student_hard_loss + pdn_student_penalty_loss

            autoencoder_student_loss = torch.mean(
                (autoencoder_result - student_result[:, channels_teacher:, :, :]) ** 2
            )
            total_loss = pdn_student_loss + autoencoder_student_loss

            total_batch = batch + epoch * len(dataloader_good)
            print(f"batch: {total_batch}/{epochs * len(dataloader_good)}  loss: {total_loss.item()}")
            if tensorboard_writer is not None:
                tensorboard_writer.add_scalar("student training", total_loss.item(), total_batch)

            optimizer.zero_grad()
            total_loss.backward()
            optimizer.step()
            scheduler.step()

        torch.save(student_pdn, "models/tmp/student.pt")

    torch.save(student_pdn, "models/student.pt")
    print("finished training student!")

    return student_pdn


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-d", "--device", choices=["cpu", "cuda"])
    parser.add_argument("--skip-teacher", action="store_true", default=False)
    parser.add_argument("--skip-autoencoder", action="store_true", default=False)
    parser.add_argument("--skip-student", action="store_true", default=False)
    parser.add_argument("--skip-quantiles", action="store_true", default=False)
    parser.add_argument("--model-config", action="store", default="model_config.yaml")
    args = parser.parse_args()
    if args.device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
    else:
        device = args.device
    print(f"using {device}")

    tensorboard_writer = SummaryWriter()

    torch.manual_seed(1337)

    with open(args.model_config) as config_file:
        model_config = yaml.safe_load(config_file)

    # if the dataset is gated/private, make sure you have run huggingface-cli login
    try:
        generic_dataset = TensorConvertedIterableDataset(
            ConvertedHuggingFaceIterableDataset(
                datasets.load_dataset("imagenet-1k", trust_remote_code=True, streaming=True)["train"]
            )
        )
    except datasets.exceptions.DatasetNotFoundError:
        raise Exception("huggingface missing token. run 'huggingface-cli login'")
    good_dataset = TensorConvertedIterableDataset(
        MVTecLOCOIterableDataset(dataset_name="mvtec_loco", group="splicing_connectors", phase="train", sorting="good")
    )

    if args.skip_teacher:
        teacher_pdn = torch.load(model_config["teacher_path"], map_location=device)
    else:
        teacher_pdn = train_teacher(
            channels=model_config["out_channels"]["teacher"],
            generic_dataset=generic_dataset,
            good_dataset=good_dataset,
            device=device,
            tensorboard_writer=tensorboard_writer,
            wideresnet_feature_layer_index=2,
        )

    if args.skip_autoencoder:
        autoencoder = torch.load(model_config["autoencoder_path"], map_location=device)
    else:
        autoencoder = train_autoencoder(
            channels=model_config["out_channels"]["autoencoder"],
            output_size=(
                model_config["out_resolution"]["width"],
                model_config["out_resolution"]["height"],
            ),
            dataset=good_dataset,
            device=device,
            teacher=teacher_pdn,
            tensorboard_writer=tensorboard_writer,
        )

    if args.skip_student:
        student_pdn = torch.load(model_config["student_path"], map_location=device)
    else:
        student_pdn = train_student(
            channels_teacher=model_config["out_channels"]["teacher"],
            channels_autoencoder=model_config["out_channels"]["autoencoder"],
            good_dataset=good_dataset,
            generic_dataset=generic_dataset,
            device=device,
            teacher=teacher_pdn,
            autoencoder=autoencoder,
            tensorboard_writer=tensorboard_writer,
        )

    if not args.skip_quantiles:
        find_quantiles(
            teacher=teacher_pdn,
            autoencoder=autoencoder,
            student=student_pdn,
            teacher_channels=model_config["out_channels"]["teacher"],
            dataset=good_dataset,
            device=device,
        )

    tensorboard_writer.close()


if __name__ == "__main__":
    main()