train.py

import os
import numpy as np
import torch
import torchvision
import argparse
from modules import resnet, network, losses
from utils import yaml_config_hook, save_model
from torch.utils import data

import random
import datetime

import torch.nn as nn
import torch.nn.functional as F
from torch.nn.functional import cosine_similarity
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from datasets.SVI import TrainDataset, TestDataset

from cluster import inference
from evaluation.evaluation import evaluate, evaluate_overclustering

from collections import OrderedDict
from sklearn.metrics import f1_score


def load_model(model, checkpoint, is_multi_gpu):
    state_dict = checkpoint['net']
    new_state_dict = OrderedDict()

    if is_multi_gpu:
        for k, v in state_dict.items():
            if k.startswith('module.'):
                new_state_dict[k] = v  
            else:
                new_state_dict['module.' + k] = v
    else:
        for k, v in state_dict.items():
            if k.startswith('module.'):
                new_state_dict[k[7:]] = v  
            else:
                new_state_dict[k] = v  

    model.load_state_dict(new_state_dict)
    return model


def get_data_loader(mode, batch_size, positive_sample_num=None, num_workers=4):
    
    if mode == 'train':
        csv_path = "./datasets/svi_dataset_example.csv"

        train_dataset = TrainDataset(csv_path, top_k=positive_sample_num)
        dataloader = DataLoader(train_dataset, batch_size=batch_size, num_workers=num_workers, pin_memory=True, shuffle=True, drop_last=False)

    elif mode == 'test':
        csv_path = "./dataset/test_dataset_example.csv"
        test_dataset = TestDataset(csv_path)
        dataloader = DataLoader(test_dataset, batch_size=batch_size, num_workers=num_workers, pin_memory=True, shuffle=False)

    else:
        print("mode input is invalid", flush=True)
        
    return dataloader


def train(train_dl, cluster_num, device):

    loss_epoch = 0
    train_interval = max(1, len(train_dl) // 20)
    validation_interval = max(1, len(train_dl) // 10)

    for step, (anchors, neighbors) in enumerate(train_dl):

        optimizer.zero_grad()
        
        B, anchor_number, C, H, W = anchors.size()
        _, neighbor_number, _, _, _ = neighbors.size()

        anchors = anchors.to(device).view(-1, C, H, W)  
        neighbors = neighbors.to(device).view(-1, C, H, W) 
        
        anchors_z, anchors_c = model(anchors) 
        neighbors_z, neighbors_c = model(neighbors)

        anchors_z = anchors_z.view(B, anchor_number, -1) 
        anchors_c = anchors_c.view(B, anchor_number, -1) 
        neighbors_z = neighbors_z.view(B, neighbor_number, -1)
        neighbors_c = neighbors_c.view(B, neighbor_number, -1) 

        anchors_c_i = anchors_c[:, 0, :] 
        anchors_c_j = anchors_c[:, 1, :] 
        neighbors_c_i = neighbors_c[:, 0, :] 
        neighbors_c_j = neighbors_c[:, 1, :]

        loss_instance = criterion_instance(anchors_z, neighbors_z)
  
        loss_cluster_anchor, ne_loss_anchor = criterion_cluster(anchors_c_i, neighbors_c_i)  
        loss_cluster_neighbor, ne_loss_neighbor = criterion_cluster(anchors_c_j, neighbors_c_j)

        loss_cluster = loss_cluster_anchor + loss_cluster_neighbor

        loss_ne = ne_loss_anchor + ne_loss_neighbor

        loss = loss_instance + loss_cluster + 0.2*loss_ne

        loss.backward()
        optimizer.step()

        if step % train_interval == 0:
            print(f"Step [{step}/{len(train_dl)}]\t loss_instance: {loss_instance.item()}\t loss_cluster: {loss_cluster.item()}", flush=True)

        if step % validation_interval == 0:
            model.eval()

            X, Y = inference(test_dl, model, device)

            if cluster_num == 5:
                nmi, ari, f, acc = evaluate(Y, X)
            else:
                nmi, ari, f, acc = evaluate_overclustering(Y, X)
                
            f1_scores = f1_score(Y, X, average=None, labels=[0, 1, 2, 3, 4])
            mean_per_class_f1_score = f1_scores.mean()

            print('NMI = {:.4f} ARI = {:.4f} F = {:.4f} ACC = {:.4f} mF1 = {:.4f}'.format(nmi, ari, f, acc, mean_per_class_f1_score), flush=True)

            model.train()
        
        loss_epoch += loss.item()

    return loss_epoch


if __name__ == "__main__":

    current_time = datetime.datetime.now()

    print("time：", current_time, flush=True)

    parser = argparse.ArgumentParser()
    config = yaml_config_hook("config/config.yaml")

    for k, v in config.items():
        parser.add_argument(f"--{k}", default=v, type=type(v))
    args = parser.parse_args()

    exp_name = str(args.exp_name)
    os.makedirs("./"+exp_name, exist_ok=False)

    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    torch.cuda.manual_seed(args.seed)
    np.random.seed(args.seed)

    print("config: ", str(config), flush=True)

    batch_size = args.batch_size
    positive_sample_num = args.positive_sample_num

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print("device: ", device, flush=True)

    train_dl = get_data_loader("train", batch_size=args.batch_size, positive_sample_num= positive_sample_num, num_workers=args.workers)
    test_dl = get_data_loader("test", batch_size=args.batch_size, num_workers=args.workers, positive_sample_num=positive_sample_num)

    cluster_num  = int(args.cluster_num)
 
    print("train_dl: ", str(len(train_dl)), flush=True)

    res = resnet.get_resnet(args.resnet)
    model = network.Network(res, args.feature_dim, cluster_num).to(device)
    print("Model loaded")

    optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, patience=5, verbose=True) 

    if torch.cuda.device_count() > 1:
        print(f"Let's use {torch.cuda.device_count()} GPUs!", flush=True)
        model = nn.DataParallel(model)

    if args.reload:
        model_fp = os.path.join(args.model_path, "checkpoint_{}.tar".format(args.start_epoch))
        checkpoint = torch.load(model_fp)
        model = load_model(model, checkpoint, torch.cuda.device_count() > 1)
        optimizer.load_state_dict(checkpoint['optimizer'])
        args.start_epoch = checkpoint['epoch'] + 1

    loss_device = device
 
    # loss functions
    criterion_instance = losses.spatial_instance_loss(args.instance_temperature, positive_sample_num, device).to(device)
    criterion_cluster = losses.spatial_cluster_loss(cluster_num, args.cluster_temperature, loss_device).to(device)

    # train
    for epoch in range(args.start_epoch, args.epochs):
        lr = str(optimizer.param_groups[0]["lr"])
        model.train()
        print(f"lr: {lr}", flush=True)

        loss_epoch = train(train_dl, cluster_num, device)

        if epoch % 5 == 0:
            save_model(args, model, optimizer, epoch)

        print(f"Epoch [{epoch}/{args.epochs}]\t Loss: {loss_epoch / len(train_dl)}", flush=True)

        scheduler.step(loss_epoch)

        torch.cuda.empty_cache()

    save_model(args, model, optimizer, args.epochs)