run_CheXpert_CM.py

import os
import argparse

import numpy as np
import torch
import torchvision.datasets as datasets
import torchvision.transforms as transforms

from utils import *
from Protopnet import ProtoPNet
from train_or_test import *
from push_prot_chex import *
from parameters import *

seed = 12
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
    
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

parser = argparse.ArgumentParser() 
parser.add_argument("--num_classes", '-nc', default=2, required=False, help="Number of classes")
parser.add_argument("--epochs", '-e', type=str, required=True, help="Number of training epochs")
parser.add_argument("--class_name", '-c', type=str, required=True, help="Name of a class")
args = parser.parse_args()

num_classes = int(args.num_classes)
shape0 = 10*num_classes
prototype_shape = (shape0, 128, 1, 1)

normalize = transforms.Normalize(mean=mean,
                                 std=std)

num_train_epoch = int(args.epochs)
num_warm_epoch = 5
push_epochs = np.linspace(10, 100, 10)[:-1]
push_start = 10

data_path = args.class_name + '/'
train_dir = data_path + 'train/'
test_dir = data_path + 'test/'
train_push_dir = data_path + 'push/'

# train set
train_dataset = datasets.ImageFolder(
    train_dir,
    transforms.Compose([
        transforms.Resize(size=(img_size, img_size)),
        transforms.ToTensor(),
        normalize,
    ]))
train_loader = torch.utils.data.DataLoader(
    train_dataset, batch_size=train_batch_size, shuffle=True,
    num_workers=2, pin_memory=False)
# push set
train_push_dataset = datasets.ImageFolder(
    train_push_dir,
    transforms.Compose([
        transforms.Resize(size=(img_size, img_size)),
        transforms.ToTensor(),
    ]))
train_push_loader = torch.utils.data.DataLoader(
    train_push_dataset, batch_size=train_push_batch_size, shuffle=False,
    num_workers=2, pin_memory=False)
# test set
test_dataset = datasets.ImageFolder(
    test_dir,
    transforms.Compose([
        transforms.Resize(size=(img_size, img_size)),
        transforms.ToTensor(),
        normalize,
    ]))
test_loader = torch.utils.data.DataLoader(
    test_dataset, batch_size=test_batch_size, shuffle=False,
    num_workers=2, pin_memory=False)

if not os.path.exists('pretrained_models'):
    os.mkdir('pretrained_models')
if not os.path.exists('ppnet_chest'):
    os.mkdir('ppnet_chest')
if not os.path.exists('prot_chest'):
    os.mkdir('prot_chest')

if __name__ == '__main__':
    # define a model
    ppnet = ProtoPNet.construct_PPNet(base_architecture='densenet121',
                                pretrained=True,
                                img_size=img_size,
                                prot_shape=prototype_shape,
                                num_classes=num_classes,
                                prototype_activation_function=prototype_activation_function,
                                add_on_layers_type = 'regular')

    ppnet = ppnet.to(device)
    model = torch.nn.DataParallel(ppnet)

    joint_optimizer_specs = \
    [{'params': ppnet.features.parameters(), 'lr': joint_optimizer_lrs['features'], 'weight_decay': 1e-3}, # bias are now also being regularized
    {'params': ppnet.add_on_layers.parameters(), 'lr': joint_optimizer_lrs['add_on_layers'], 'weight_decay': 1e-3},
    {'params': ppnet.prototype_vectors, 'lr': joint_optimizer_lrs['prototype_vectors']},
    ]
    joint_optimizer = torch.optim.Adam(joint_optimizer_specs)
    # joint_lr_scheduler = torch.optim.lr_scheduler.StepLR(joint_optimizer, step_size=joint_lr_step_size, gamma=0.1)

    warm_optimizer_specs = \
    [{'params': ppnet.add_on_layers.parameters(), 'lr': warm_optimizer_lrs['add_on_layers'], 'weight_decay': 1e-3},
    {'params': ppnet.prototype_vectors, 'lr': warm_optimizer_lrs['prototype_vectors']},
    ]
    warm_optimizer = torch.optim.Adam(warm_optimizer_specs)

    last_layer_optimizer_specs = [{'params': ppnet.last_layer.parameters(), 'lr': last_layer_optimizer_lr}]
    last_layer_optimizer = torch.optim.Adam(last_layer_optimizer_specs)

    # Training
    for epoch in range(num_train_epoch):
        print('epoch', epoch)

        if epoch < num_warm_epoch:
            mode(model, warm=True)
            model.train()
            correct_ratio, loss = train_or_test(model, train_loader, warm_optimizer, class_specific=True)

        else:
            mode(model, joint=True)
            model.train()
            # joint_lr_scheduler.step()
            _, loss = train_or_test(model, train_loader, joint_optimizer, class_specific=True)

        model.eval()
        print('---test---')
        acc, loss_test = train_or_test(model, test_loader, class_specific=True)
        save_model_w_condition(model=ppnet, model_dir=model_dir, model_name=str(epoch) + 'nopush', acc=acc, target_acc=0.60)

        if epoch >= push_start and epoch in push_epochs:
            push_prototypes(
                train_push_loader, # pytorch dataloader (must be unnormalized in [0,1])
                prototype_network_parallel=model, # pytorch network with prototype_vectors
                class_specific=True,
                preprocess_input_function=preprocess_input_function, # normalize if needed
                prototype_layer_stride=1,
                root_dir_for_saving_prototypes=root_dir_for_saving_prototypes, # if not None, prototypes will be saved here
                epoch_number=epoch, # if not provided, prototypes saved previously will be overwritten
                prototype_img_filename_prefix=prototype_img_filename_prefix,
                prototype_self_act_filename_prefix=prototype_self_act_filename_prefix,
                proto_bound_boxes_filename_prefix=proto_bound_boxes_filename_prefix,
                save_prototype_class_identity=True)

            model.eval()
            acc, loss_test = train_or_test(model, test_loader, class_specific=True)

            if prototype_activation_function != 'linear':
                mode(model, last=True)
                for i in range(12):
                    print('iteration: \t{0}'.format(i))
                    _, loss = train_or_test(model, train_loader, last_layer_optimizer, class_specific=True)
                    acc, loss_test = train_or_test(model, test_loader, class_specific=True)
                    save_model_w_condition(model=ppnet, model_dir=model_dir, model_name=str(epoch) + '_' + str(i) + 'push', acc=acc, target_acc=0.60)