train_dne_gezo.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
# --------------------------------------------------------
# We changed it for adapted in our paper.
# --------------------------------------------------------

import argparse
import datetime
import json
import numpy as np
import os
import time
from pathlib import Path
import copy

import torch
import torch.backends.cudnn as cudnn
from torch.utils.tensorboard import SummaryWriter

import timm

# assert timm.__version__ == "0.3.2" # version check
from timm.models.layers import trunc_normal_
from timm.data.mixup import Mixup

import util.misc as misc
from util.datasets import build_dataset_chest_xray
from util.pos_embed import interpolate_pos_embed
import models_vit

from engine_finetune import train_dne_gezo
from engine_finetune import DNELayer
from pprint import pprint

def get_args_parser():
    parser = argparse.ArgumentParser('MAE fine-tuning for image classification', add_help=False)
    parser.add_argument('--batch_size', default=32, type=int,
                        help='Batch size per GPU (effective batch size is batch_size * accum_iter * # gpus')
    parser.add_argument('--epochs', default=50, type=int)
    parser.add_argument('--accum_iter', default=1, type=int,
                        help='Accumulate gradient iterations (for increasing the effective batch size under memory constraints)')

    # Model parameters
    parser.add_argument('--model', default='vit_large_patch16', type=str, metavar='MODEL',
                        help='Name of model to train')

    parser.add_argument('--input_size', default=224, type=int,
                        help='images input size')

    parser.add_argument('--drop_path', type=float, default=0.1, metavar='PCT',
                        help='Drop path rate (default: 0.1)')

    # Optimizer parameters
    parser.add_argument('--clip_grad', type=float, default=None, metavar='NORM',
                        help='Clip gradient norm (default: None, no clipping)')

    # Augmentation parameters
    parser.add_argument('--smoothing', type=float, default=0.1,
                        help='Label smoothing (default: 0.1)')

    # * Mixup params
    parser.add_argument('--mixup', type=float, default=0,
                        help='mixup alpha, mixup enabled if > 0.')
    parser.add_argument('--cutmix', type=float, default=0,
                        help='cutmix alpha, cutmix enabled if > 0.')
    parser.add_argument('--cutmix_minmax', type=float, nargs='+', default=None,
                        help='cutmix min/max ratio, overrides alpha and enables cutmix if set (default: None)')
    parser.add_argument('--mixup_prob', type=float, default=1.0,
                        help='Probability of performing mixup or cutmix when either/both is enabled')
    parser.add_argument('--mixup_switch_prob', type=float, default=0.5,
                        help='Probability of switching to cutmix when both mixup and cutmix enabled')
    parser.add_argument('--mixup_mode', type=str, default='batch',
                        help='How to apply mixup/cutmix params. Per "batch", "pair", or "elem"')

    # * Finetuning params
    parser.add_argument('--finetune', default='',
                        help='finetune from checkpoint')
    parser.add_argument('--global_pool', action='store_true')
    parser.set_defaults(global_pool=True)

    # Dataset parameters
    parser.add_argument('--data_path', default='/datasets01/imagenet_full_size/061417/', type=str,
                        help='dataset path')
    parser.add_argument('--nb_classes', default=1000, type=int,
                        help='number of the classification types')

    parser.add_argument('--output_dir', default='./output_dir',
                        help='path where to save, empty for no saving')
    parser.add_argument('--log_dir', default='./output_dir',
                        help='path where to tensorboard log')
    parser.add_argument('--device', default='cuda',
                        help='device to use for training / testing')
    parser.add_argument('--seed', default=0, type=int)

    parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
                        help='start epoch')
    parser.add_argument('--eval', action='store_true',
                        help='Perform evaluation only')
    parser.add_argument('--dist_eval', action='store_true', default=False,
                        help='Enabling distributed evaluation (recommended during training for faster monitor')
    parser.add_argument('--num_workers', default=10, type=int)
    parser.add_argument('--pin_mem', action='store_true',
                        help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
    parser.add_argument('--no_pin_mem', action='store_false', dest='pin_mem')
    parser.set_defaults(pin_mem=True)

    # distributed training parameters
    parser.add_argument('--world_size', default=1, type=int,
                        help='number of distributed processes')
    parser.add_argument('--local_rank', default=-1, type=int)
    parser.add_argument('--dist_on_itp', action='store_true')
    parser.add_argument('--dist_url', default='env://',
                        help='url used to set up distributed training')
    parser.add_argument('--eval_interval', default=2, type=int)
    parser.add_argument('--vit_dropout_rate', type=float, default=0.1,
                        help='Dropout rate for ViT blocks (default: 0.1)')
    parser.add_argument("--dataset", default='chestxray', type=str)

    parser.add_argument('--loss_func', default=None, type=str)
    
    parser.add_argument('--sa_path', default="./sa_cls.pt", type=str, help="the path to the sensitive attribute classifier")

    parser.add_argument("--csv_path", default=None, type=str)

    parser.add_argument("--disease", default=None, type=str)

    parser.add_argument("--lambda_1", default=1.0, type=float)

    parser.add_argument("--lambda_reg", default=0.01, type=float)

    parser.add_argument('--dne', action='store_true')

    parser.add_argument("--init_step", default=0.01, type=float)

    parser.add_argument("--sampled_steps", default=10, type=int)

    parser.add_argument("--momentum", default=0.9, type=int)

    return parser


def main(args):
    misc.init_distributed_mode(args)

    print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__))))
    print("{}".format(args).replace(', ', ',\n'))

    device = torch.device(args.device)

    # fix the seed for reproducibility
    seed = args.seed + misc.get_rank()
    torch.manual_seed(seed)
    np.random.seed(seed)

    cudnn.benchmark = True

    dataset_train = build_dataset_chest_xray(split='train', args=args)
    dataset_test = build_dataset_chest_xray(split='test', args=args)

    num_tasks = misc.get_world_size()
    global_rank = misc.get_rank()


    
    sampler_train = torch.utils.data.DistributedSampler(
        dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
    )
    print("Sampler_train = %s" % str(sampler_train))

    if args.dist_eval:
        if len(dataset_test) % num_tasks != 0:
            print('Warning: Enabling distributed evaluation with an eval dataset not divisible by process number. '
                    'This will slightly alter validation results as extra duplicate entries are added to achieve '
                    'equal num of samples per-process.')
        # sampler_val = torch.utils.data.DistributedSampler(
        #     dataset_val, num_replicas=num_tasks, rank=global_rank, shuffle=True)  # shuffle=True to reduce monitor bias
        sampler_test = torch.utils.data.DistributedSampler(
            dataset_test, num_replicas=num_tasks, rank=global_rank, shuffle=True)
    else:
        sampler_test = torch.utils.data.SequentialSampler(dataset_test)

    if global_rank == 0 and args.log_dir is not None and not args.eval:
        os.makedirs(args.log_dir, exist_ok=True)
        log_writer = SummaryWriter(log_dir=args.log_dir)
    else:
        log_writer = None

    data_loader_train = torch.utils.data.DataLoader(
        dataset_train, sampler=sampler_train,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        pin_memory=args.pin_mem,
        drop_last=True,
    )

    mixup_fn = None
    mixup_active = args.mixup > 0 or args.cutmix > 0. or args.cutmix_minmax is not None
    if mixup_active:
        print("Mixup is activated!")
        mixup_fn = Mixup(
            mixup_alpha=args.mixup, cutmix_alpha=args.cutmix, cutmix_minmax=args.cutmix_minmax,
            prob=args.mixup_prob, switch_prob=args.mixup_switch_prob, mode=args.mixup_mode,
            label_smoothing=args.smoothing, num_classes=args.nb_classes)
        
    if 'vit' in args.model:
        model = models_vit.__dict__[args.model](
            img_size=args.input_size,
            num_classes=args.nb_classes,
            drop_rate=args.vit_dropout_rate,
            drop_path_rate=args.drop_path,
            global_pool=args.global_pool,
        )
    else:
        raise NotImplementedError("The existing implementation only supports MedMAE. More FMs will be added here.")


    if args.finetune and not args.eval:
        if 'vit' in args.model:
            if args.finetune == "imgnet": # load model from timm pretrained supervised vit
                pretrained_model = timm.create_model('vit_base_patch16_224', pretrained=True)
                pretrained_model.eval()
                checkpoint_model = pretrained_model.state_dict()
                state_dict = model.state_dict()
                for k in ['head.weight', 'head.bias']:
                    if k in checkpoint_model and checkpoint_model[k].shape != state_dict[k].shape:
                        print(f"Removing key {k} from pretrained checkpoint")
                        del checkpoint_model[k]
                if args.global_pool:
                    for k in ['fc_norm.weight', 'fc_norm.bias']:
                        try:
                            del checkpoint_model[k]
                        except:
                            pass
            else: # load pretrained model from the given path
                checkpoint = torch.load(args.finetune, map_location='cpu')
                print("Load pre-trained checkpoint from: %s" % args.finetune)
                checkpoint_model = checkpoint['model']
                state_dict = model.state_dict()
                for k in ['head.weight', 'head.bias']:
                    if k in checkpoint_model and checkpoint_model[k].shape != state_dict[k].shape:
                        print(f"Removing key {k} from pretrained checkpoint")
                        del checkpoint_model[k]
                if args.global_pool:
                    for k in ['fc_norm.weight', 'fc_norm.bias']:
                        try:
                            del checkpoint_model[k]
                        except:
                            pass
                 

            # interpolate position embedding
            interpolate_pos_embed(model, checkpoint_model)

            # load pre-trained model
            msg = model.load_state_dict(checkpoint_model, strict=False)
            print(msg)

            print("Freeze all parameters except head")
            # Freeze all the parameters in the model
            for param in model.parameters():
                param.requires_grad = False

            # Unfreeze the parameters in the classifier head
            if hasattr(model, 'head'):
                for param in model.head.parameters():
                    param.requires_grad = True
            elif hasattr(model, 'classifier'):
                for param in model.classifier.parameters():
                    param.requires_grad = True
            else:
                raise NameError("The model does not have a recognizable classifier head")

            # manually initialize fc layer
            trunc_normal_(model.head.weight, std=2e-5)
        else:
            raise NotImplementedError("Only ViT is supported at this point. More FMs will be added here.")

    model_sa = copy.deepcopy(model) # model_sa is the SA classifier
    model_sa.to(device)
    noise_layer = DNELayer((224,224)).to(device)

    model_sa_without_ddp = model_sa

    n_parameters = sum(p.numel() for p in model_sa.parameters() if p.requires_grad)

    print("Model = %s" % str(model_sa_without_ddp))
    print('number of params (M): %.2f' % (n_parameters / 1.e6))

    eff_batch_size = args.batch_size * args.accum_iter * misc.get_world_size()
    

    print("accumulate grad iterations: %d" % args.accum_iter)
    print("effective batch size: %d" % eff_batch_size)

    if args.distributed:
        model_sa = torch.nn.parallel.DistributedDataParallel(model_sa, device_ids=[args.gpu])
        model_sa_without_ddp = model_sa.module
        noise_layer = torch.nn.parallel.DistributedDataParallel(noise_layer, device_ids=[args.gpu])
        noise_layer_without_ddp = noise_layer.module


    if args.dataset == 'chexpert':
        criterion = torch.nn.CrossEntropyLoss()
    else:
        raise NotImplementedError()


    print("criterion = %s" % str(criterion))


    # Load the SA classifier
    biased_checkpoint = torch.load("./sa_cls.pt", map_location='cpu')
    print("Load pre-trained gender classifier from: %s" % "./genderclassifier_ViT_fromMAEXray.pt")
    model_sa_state_dict = biased_checkpoint
    model_sa = copy.deepcopy(model)
    model_sa.load_state_dict(model_sa_state_dict, strict=True)
    model_sa.eval()
    model_sa.to(device)





    # Start training the DNE noise
    print(f"Start training for {args.epochs} epochs")
    start_time = time.time()
    for epoch in range(args.start_epoch, args.epochs):
        print("======================== Train DNE Epoch {} ========================".format(epoch))
        if args.distributed:
            data_loader_train.sampler.set_epoch(epoch)
        train_dne_gezo(
            model_sa,
            data_loader_train,
            device,
            epoch,
            log_writer=log_writer,
            args=args,
            dne_layer=noise_layer
        )

    torch.save(noise_layer_without_ddp.state_dict(), os.path.join(args.output_dir, "noise_layer.pt"))
    total_time = time.time() - start_time
    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
    print('Training time {}'.format(total_time_str))


if __name__ == '__main__':
    args = get_args_parser()
    args = args.parse_args()
    if args.output_dir:
        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
    main(args)