train_dne.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 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
from engine_finetune import DNELayer

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')

    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_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)
    dne_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
        dne_layer = torch.nn.parallel.DistributedDataParallel(dne_layer, device_ids=[args.gpu])
        dne_layer_without_ddp = dne_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(args.sa_path, map_location='cpu')
    print("Load pre-trained gender classifier from: %s" % args.sa_path)
    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(
            model_sa,
            data_loader_train,
            device,
            epoch,
            log_writer=log_writer,
            args=args,
            dne_layer=dne_layer
        )

    torch.save(dne_layer_without_ddp.state_dict(), os.path.join(args.output_dir, "dne_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)