Official codebase for the NeurIPS 2023 paper: Towards Last-layer Retraining for Group Robustness with Fewer Annotations. https://arxiv.org/abs/2309.08534
"""Main file for last-layer retraining experimentation."""

# Ignores nuisance warnings. Must be called first.
from milkshake.utils import ignore_warnings
ignore_warnings()

# Imports Python builtins.
from configargparse import Parser
from copy import deepcopy
from distutils.util import strtobool
import os.path as osp
import pickle
import sys

# Imports Python packages.
import numpy as np

# Imports PyTorch packages.
from pytorch_lightning import Trainer

# Imports milkshake packages.
from milkshake.args import add_input_args
from milkshake.datamodules.celeba import CelebA
from milkshake.datamodules.civilcomments import CivilComments
from milkshake.datamodules.disagreement import Disagreement
from milkshake.datamodules.multinli import MultiNLI
from milkshake.datamodules.waterbirds import Waterbirds
from milkshake.imports import valid_models_and_datamodules
from milkshake.main import main
from milkshake.utils import get_weights


class CelebADisagreement(CelebA, Disagreement):
    """DataModule for the CelebADisagreeement dataset."""

    def __init__(self, args, **kwargs):
        super().__init__(args, **kwargs)

class CivilCommentsDisagreement(CivilComments, Disagreement):
    """DataModule for the CivilCommentsDisagreement dataset."""

    def __init__(self, args, **kwargs):
        super().__init__(args, **kwargs)
        
class MultiNLIDisagreement(MultiNLI, Disagreement):
    """DataModule for the MultiNLIDisagreement dataset."""

    def __init__(self, args, **kwargs):
        super().__init__(args, **kwargs)

class WaterbirdsDisagreement(Waterbirds, Disagreement):
    """DataModule for the WaterbirdsDisagreement dataset."""

    def __init__(self, args, **kwargs):
        super().__init__(args, **kwargs)

def set_training_parameters(args):
    if args.datamodule == "waterbirds":
        args.datamodule_class = WaterbirdsDisagreement
        args.num_classes = 2
        args.retrain_epochs = 100
        args.finetune_steps = 500
        args.finetune_lrs = [1e-5, 1e-4, 1e-3, 1e-2]
    elif args.datamodule == "celeba":
        args.datamodule_class = CelebADisagreement
        args.num_classes = 2
        args.retrain_epochs = 100
        args.finetune_steps = 500
        args.finetune_lrs = [1e-5, 1e-4, 1e-3, 1e-2]
    elif args.datamodule == "civilcomments":
        args.datamodule_class = CivilCommentsDisagreement
        args.num_classes = 2
        args.retrain_epochs = 10
        args.finetune_steps = 1000
        args.finetune_lrs = [1e-7, 1e-6, 1e-5, 1e-4]
    elif args.datamodule == "multinli":
        args.datamodule_class = MultiNLIDisagreement
        args.num_classes = 3
        args.retrain_epochs = 10
        args.finetune_steps = 1000
        args.finetune_lrs = [1e-7, 1e-6, 1e-5, 1e-4]
    else:
        raise ValueError(f"DataModule {args.datamodule} not supported.")

    args.finetune_num_datas = [10, 20, 50, 100, 200, 500]
    args.dropout_probs = [0.5, 0.7, 0.9]
    args.early_stop_nums = [1, 2, 3]

def load_erm():
    if osp.isfile("erm.pkl"):
        with open("erm.pkl", "rb") as f:
            erm = pickle.load(f)
    else: 
        datasets = ["waterbirds", "celeba", "civilcomments", "multinli"]
        seeds = [1, 2, 3]
        balance = [True, False]
        split = ["train", "combined"]
        train_pct = [80, 85, 90, 95, 100]

        erm = {}
        for d in datasets:
            erm[d] = {}
            for s in seeds:
                erm[d][s] = {}
                for b in balance:
                    erm[d][s][b] = {}
                    for p in split:
                        erm[d][s][b][p] = {}
                        for t in train_pct:
                            erm[d][s][b][p][t] = {"version": -1, "metrics": []}

        with open("erm.pkl", "wb") as f:
            pickle.dump(erm, f)

    return erm

def dump_erm(args, curr_erm):
    erm = load_erm()
    erm[args.datamodule][args.seed][args.balance_erm][args.split][args.train_pct] = curr_erm

    with open("erm.pkl", "wb") as f:
        pickle.dump(erm, f)

def reset_fc_hook(model):
    try:
        for layer in model.model.fc:
            if hasattr(layer, "reset_parameters"):
                layer.reset_parameters()
    except:
        model.model.fc.reset_parameters()

def print_metrics(metrics):
    val_metrics, test_metrics = metrics

    val_group_accs = [
        acc for name, acc in val_metrics[0].items() if "group" in name
    ]
    test_group_accs = [
        acc for name, acc in test_metrics[0].items() if "group" in name
    ]

    val_avg_acc = val_metrics[0]["val_acc"]
    test_avg_acc = test_metrics[0]["test_acc"]

    val_worst_group_acc = min(val_group_accs) 
    test_worst_group_acc = min(test_group_accs) 

    val_avg_acc = round(val_avg_acc * 100, 1)
    val_worst_group_acc = round(val_worst_group_acc * 100, 1)

    test_avg_acc = round(test_avg_acc * 100, 1)
    test_worst_group_acc = round(test_worst_group_acc * 100, 1)

    print(f"Val Average Acc: {val_avg_acc}")
    print(f"Val Worst Group Acc: {val_worst_group_acc}")
    print(f"Test Average Acc: {test_avg_acc}")
    print(f"Test Worst Group Acc: {test_worst_group_acc}")
    print()

def print_results(erm_metrics, results, keys):
    print("---Experiment Results---")
    print("\nERM")
    print_metrics(erm_metrics)

    for key in keys:
        print(key.title())
        if "self" in key:
            print(f"Best params: {results[key]['params']}")
        print_metrics(results[key]["metrics"])

def finetune_last_layer(
    args,
    finetune_type,
    model_class,
    dropout_prob=0,
    early_stop_weights=None,
    finetune_num_data=None,
    worst_group_pct=None,
):
    disagreement_args = deepcopy(args)
    disagreement_args.finetune_type = finetune_type
    disagreement_args.dropout_prob = dropout_prob

    finetune_args = deepcopy(args)
    finetune_args.train_fc_only = True
    finetune_args.lr = args.finetune_lr
    finetune_args.lr_scheduler = "step"
    finetune_args.lr_steps = []

    # Sets parameters for finetuning (first) or retraining (second)
    if finetune_num_data:
        finetune_args.max_steps = args.finetune_steps
        reset_fc = False
    else:
        finetune_args.max_epochs = args.retrain_epochs
        reset_fc = True

    if finetune_args.save_self_model:
        finetune_args.check_val_every_n_steps = args.finetune_steps
        finetune_args.ckpt_every_n_steps = args.finetune_steps
    else:
        finetune_args.check_val_every_n_steps = args.finetune_steps + 1
        finetune_args.ckpt_every_n_steps = args.finetune_steps + 1


    model = model_class(disagreement_args)
    early_stop_model = None
    if early_stop_weights:
        early_args = deepcopy(disagreement_args)
        early_args.weights = early_stop_weights
        early_stop_model = model_class(early_args)

    def disagreement_class(orig_datamodule_class):
        class Disagreement(orig_datamodule_class):
            def __init__(self, args):
                super().__init__(
                    disagreement_args,
                    early_stop_model=early_stop_model,
                    model=model,
                    num_data=finetune_num_data,
                    worst_group_pct=worst_group_pct,
                )

        return Disagreement

    model_hooks = [reset_fc_hook] if reset_fc else None
    _, val_metrics, test_metrics = main(
        finetune_args,
        model_class,
        disagreement_class(args.datamodule_class),
        model_hooks=model_hooks,
    )

    return val_metrics, test_metrics

def experiment(args, model_class):
    # Loads ERM paths and metrics from pickle file.
    erm = load_erm()
    
    # Adds experiment-specific parameters to args.
    set_training_parameters(args)

    # Trains ERM model.
    curr_erm = erm[args.datamodule][args.seed][args.balance_erm][args.split][args.train_pct]
    erm_version = curr_erm["version"]
    erm_metrics = curr_erm["metrics"]
    if erm_version == -1:
        args.balanced_sampler = True if args.balance_erm else False
        model, erm_val_metrics, erm_test_metrics = main(args, model_class, args.datamodule_class)
        args.balanced_sampler = False

        erm_version = model.trainer.logger.version
        erm_metrics = [erm_val_metrics, erm_test_metrics]

        curr_erm["version"] = erm_version
        curr_erm["metrics"] = erm_metrics
        dump_erm(args, curr_erm)
        del model

    elif not erm_metrics:
        args.weights = get_weights(args, erm_version, idx=-1)
        args.eval_only = True
        _, erm_val_metrics, erm_test_metrics = main(args, model_class, args.datamodule_class)
        args.eval_only = False

        erm_metrics = [erm_val_metrics, erm_test_metrics]
        curr_erm["metrics"] = erm_metrics
        dump_erm(args, curr_erm)

    def print_results2(results, keys):
        return print_results(erm_metrics, results, keys)

    print("ERM")
    print_metrics(erm_metrics)

    # When these two arguments are passed, the entire held-out set
    # (except the actual validation set) is used for training. So,
    # there is no data left for last-layer retraining, and we return.
    if args.split == "combined" and args.train_pct == 100:
        return

    # Gets last-epoch ERM weights.
    args.weights = get_weights(args, erm_version, idx=-1)

    # Sets finetune types. Note that "group-unbalanced retraining" will be
    # either class-unbalanced or class-balanced based on the value
    # of args.balanced_sampler. Note that this does NOT set which techniques
    # will be run during the experiment (it just sets the strings). To run only
    # specific techniques, comment out the calls below or use --no_self.
    finetune_types = [
        "group-unbalanced retraining",
        "group-balanced retraining",
        "misclassification self",
        "early-stop misclassification self",
        "dropout disagreement self",
        "early-stop disagreement self",
    ]

    # Prepares results dictionary.
    results = {f: {"val_worst_group_acc": -1, "metrics": [], "params": []}
               for f in finetune_types}

    def finetune_helper(
        finetune_type,
        dropout_prob=0,
        early_stop_num=None,
        finetune_lr=None,
        finetune_num_data=None,
        worst_group_pct=None,
    ):
        args.finetune_lr = finetune_lr if finetune_lr else args.lr

        param_str = f" LR {args.finetune_lr}"
        if finetune_num_data:
            param_str += f" Num Data {finetune_num_data}"
        if dropout_prob:
            param_str += f" Dropout {dropout_prob}"
        if early_stop_num:
            param_str += f" Early Stop {early_stop_num}"
        if worst_group_pct:
            param_str += f" Worst-group Pct {worst_group_pct}"
        print(f"{finetune_type.title()}{param_str}")

        early_stop_weights = None
        if early_stop_num:
            early_stop_weights = get_weights(args, erm_version, idx=early_stop_num-1)
        val_metrics, test_metrics = finetune_last_layer(
            args,
            finetune_type,
            model_class,
            dropout_prob=dropout_prob,
            early_stop_weights=early_stop_weights,
            finetune_num_data=finetune_num_data,
            worst_group_pct=worst_group_pct,
        )

        print_metrics([val_metrics, test_metrics])

        val_worst_group_acc = min([acc for name, acc in val_metrics[0].items() if "group" in name])

        if val_worst_group_acc >= results[finetune_type]["val_worst_group_acc"]:
            results[finetune_type]["val_worst_group_acc"] = val_worst_group_acc
            results[finetune_type]["metrics"] = [val_metrics, test_metrics]

            params = []
            if finetune_num_data:
                params.append(finetune_num_data)
            if finetune_lr:
                params.append(finetune_lr)
            if dropout_prob:
                params.append(dropout_prob)
            if early_stop_num:
                params.append(early_stop_num)
            results[finetune_type]["params"] = params

    # Performs worst-group data ablation.
    if args.worst_group_ablation:
        for worst_group_pct in [2.5, 5] + [12.5 * j for j in range(1, 9)]:
            finetune_helper("class-balanced retraining", worst_group_pct=worst_group_pct)
        return

    # Performs early-stop disagreement demo.
    if args.demo:
        finetune_helper(
            "early-stop disagreement self",
            early_stop_num=1,
            finetune_lr=1e-2,
            finetune_num_data=100,
        )
        print_results2(results, finetune_types[-1:])
        return

    # Performs last-layer retraining.
    finetune_helper("group-unbalanced retraining")
    finetune_helper("group-balanced retraining")

    if args.no_self:
        print_results2(results, finetune_types[:2])
        return

    # Perform SELF hyperparameter search using worst-group validation accuracy.
    for finetune_num_data in args.finetune_num_datas:
        for finetune_lr in args.finetune_lrs:
            def finetune_helper2(finetune_type, dropout_prob=0, early_stop_num=None):
                finetune_helper(
                    finetune_type,
                    dropout_prob=dropout_prob,
                    early_stop_num=early_stop_num,
                    finetune_lr=finetune_lr,
                    finetune_num_data=finetune_num_data,
                )

            finetune_helper2("misclassification self")

            for early_stop_num in args.early_stop_nums:
                finetune_helper2("early-stop misclassification self", early_stop_num=early_stop_num)

            for dropout_prob in args.dropout_probs:
                finetune_helper2("dropout disagreement self", dropout_prob=dropout_prob)

            
            for early_stop_num in args.early_stop_nums:
                finetune_helper2("early-stop disagreement self", early_stop_num=early_stop_num)

    print_results2(results, finetune_types)

    
if __name__ == "__main__":
    parser = Parser(
        args_for_setting_config_path=["-c", "--cfg", "--config"],
        config_arg_is_required=True,
    )

    parser = add_input_args(parser)
    parser = Trainer.add_argparse_args(parser)
    
    parser.add("--balance_erm", default=True, type=lambda x: bool(strtobool(x)),
               help="Whether to perform class-balancing during ERM training.")
    parser.add("--balance_finetune", choices=["sampler", "subset", "none"], default="sampler",
               help="Which type of class-balancing to perform during finetuning.")
    parser.add("--demo", action="store_true",
               help="Whether to run a quick demo of early-stop disagreement SELF.")
    parser.add("--no_self", action="store_true",
               help="Whether to perform only ERM and last-layer retraining, no SELF.")
    parser.add("--save_self_model", action="store_true",
               help="Whether to save the SELF model outputs.")
    parser.add("--split", choices=["combined", "train"], default="train",
               help="The split to train on; either the train set or the combined train and held-out set.")
    parser.add("--train_pct", default=100, type=int,
               help="The percentage of the train set to utilize (for ablations)")
    parser.add("--val_pct", default=100, type=int,
               help="The percentage of the val set to utilize (for ablations)")
    parser.add("--worst_group_ablation", action="store_true",
               help="Whether to perform an ablation on the amount of worst-group data used during retraining.")

    args = parser.parse_args()

    models, _ = valid_models_and_datamodules()
        
    experiment(args, models[args.model])


exps/disagreement.py
tmlabonte/last-layer-retraining
"""Defines configuration parameters for models and datamodules."""

# Imports Python packages.
from configargparse import Parser
from distutils.util import strtobool

# Import PyTorch packages.
from pytorch_lightning import Trainer

# Imports milkshake packages.
from milkshake.imports import valid_model_and_datamodule_names


def parse_args():
    """Reads configuration file and returns configuration dictionary."""

    parser = Parser(
        args_for_setting_config_path=["-c", "--cfg", "--config"],
        config_arg_is_required=True,
    )

    parser = add_input_args(parser)
    parser = Trainer.add_argparse_args(parser)
    parser.set_defaults(strategy="ddp_find_unused_parameters_false")

    args = parser.parse_args()

    return args

def add_input_args(parser):
    """Loads configuration parameters into given configargparse.Parser."""

    model_names, datamodule_names = valid_model_and_datamodule_names()

    parser.add("--balanced_sampler", default=False, type=lambda x: bool(strtobool(x)),
               help="Whether to use a class-balanced random sampler for training.")
    parser.add("--batch_size", type=int,
               help="The number of samples to include per batch.")
    parser.add("--bias", default=True, type=lambda x: bool(strtobool(x)),
               help="Whether to add a bias parameter to each model weight.")
    parser.add("--check_val_every_n_epochs", default=1, dest="check_val_every_n_epoch", type=int,
               help="The number of epochs after which to run validation.")
    parser.add("--check_val_every_n_steps", dest="val_check_interval", type=int,
               help="The number of steps after which to run validation.")
    parser.add("--ckpt_every_n_epochs", default=1, type=int,
               help="The number of epochs after which a model checkpoint will be saved.")
    parser.add("--ckpt_every_n_steps", type=int,
               help="The number of steps after which a model checkpoint will be saved.")
    parser.add("--class_weights", default=[], nargs="*", type=float,
               help="The amount by which to weight each class in the loss function.")
    parser.add("--data_augmentation", default=True, type=lambda x: bool(strtobool(x)),
               help="Whether to use data augmentation as specified in the DataModule.")
    parser.add("--data_dir", default="data",
               help="The name of the directory where data will be saved.")
    parser.add("--datamodule", choices=datamodule_names,
               help="The name of the DataModule to utilize.")
    parser.add("--dropout_prob", default=0, type=float,
               help="The probability by which a node will be dropped out.")
    parser.add("--eval_only", default=False, type=lambda x: bool(strtobool(x)),
               help="Whether to skip training and only evaluate the model.")
    parser.add("--input_channels", default=3, type=int,
               help="The number of channels for image inputs.")
    parser.add("--label_noise", default=0, type=float,
               help="The probability by which class labels will be flipped.")
    parser.add("--loss", choices=["cross_entropy", "mse"], default="cross_entropy",
               help="The name of the loss function to utilize for optimization.")
    parser.add("--lr", type=float,
               help="The learning rate to utilize for optimization.") 
    parser.add("--lr_drop", default=0.1, type=float,
               help="The factor by which to drop the LR when using the step scheduler.")
    parser.add("--lr_scheduler", choices=["cosine", "cosine_warmup", "linear", "step"], default="step",
               help="The name of the LR scheduler to utilize.")
    parser.add("--lr_steps", default=[], nargs="*", type=int,
               help="A list of epochs to drop the LR when using the step scheduler.")
    parser.add("--lr_warmup_epochs", default=0, type=int,
               help="The number of epochs to warm up using certain schedulers.")
    parser.add("--metrics", nargs="+", default=["loss", "acc", "acc_by_class"],
               help=("Which metrics to compute and log. Choices include loss, acc, acc5, acc_by_class"
                     "acc5_by_class, and any custom metrics implemented in a DataModule."))
    parser.add("--model", choices=model_names,
               help="The name of the Model to utilize.")
    parser.add("--momentum", default=0.9, type=float,
               help="The momentum value to utilize with the SGD optimizer.")
    parser.add("--num_workers", default=4, type=int,
               help="The number of sub-processes to use for data loading.")
    parser.add("--optimizer", choices=["adam", "adamw", "sgd"], default="sgd",
               help="The name of the optimizer to utilize.")
    parser.add("--out_dir", default="out",
               help="The name of the directory where outputs will be saved.")
    parser.add("--persistent_workers", default=False, type=lambda x: bool(strtobool(x)),
               help="Whether to use persistent workers (typically for ddp_spawn).")
    parser.add("--refresh_rate", default=1, type=int,
               help="The number of batches after which to update the progress bar.")
    parser.add("--resnet_small_input", default=False, type=lambda x: bool(strtobool(x)),
               help="Whether to reduce the size of the first ResNet layer for small inputs.")
    parser.add("--resnet_l1_regularization", default=0, type=float,
               help="Whether to apply l1-norm regularization to the ResNet last layer.")
    parser.add("--resnet_pretrained", default=True, type=lambda x: bool(strtobool(x)),
               help="Whether to use pretrained weights in the ResNet.")
    parser.add("--resnet_version", choices=[18, 34, 50, 101, 152], default=18, type=int,
               help="The ResNet version to utilize.")
    parser.add("--resume_training", default=False, type=lambda x: bool(strtobool(x)),
               help="Whether to resume the training state from the given checkpoint.")
    parser.add("--resume_weights", default=False, type=lambda x: bool(strtobool(x)),
               help="Whether to load the weights from the given checkpoint.")
    parser.add("--seed", default=1, type=int,
               help="The random seed to utilize.")
    parser.add("--train_fc_only", default=False, type=lambda x: bool(strtobool(x)),
               help="Whether to freeze all parameters except the last layer.")
    parser.add("--val_split", default=0.2, type=float,
               help="The proportion of training data to reserve for validation.")
    parser.add("--wandb", default=True, type=lambda x: bool(strtobool(x)),
               help="Whether to log with Weights and Biases (otherwise uses TensorBoard).")
    parser.add("--wandb_dir", default="wandb",
               help="The name of the directory where wandb outputs will be saved.")
    parser.add("--weights", default="",
               help="The filepath of the model checkpoint to load or resume.")
    parser.add("--weight_decay", default=1e-4, type=float,
               help="The l2-norm regularization to utilize during optimization.")
    
    return parser

milkshake/args.py
tmlabonte/last-layer-retraining
"""Dataset and DataModule for the CelebA dataset."""

# Imports Python builtins.
import os.path as osp

# Imports Python packages.
import numpy as np
import pandas as pd
from PIL import Image
import wget

# Imports PyTorch packages.
from torchvision.datasets.utils import (
    download_file_from_google_drive,
    extract_archive,
)
from torchvision.transforms import (
    CenterCrop,
    Compose,
    Normalize,
    RandomHorizontalFlip,
    RandomResizedCrop,
    Resize,
    ToTensor,
)

# Imports milkshake packages.
from milkshake.datamodules.dataset import Dataset
from milkshake.datamodules.datamodule import DataModule


class CelebADataset(Dataset):
    """Dataset for the CelebA dataset."""

    def __init__(self, *xargs, **kwargs):
        super().__init__(*xargs, **kwargs)

    def download(self):
        celeba_dir = osp.join(self.root, "celeba")
        if not osp.isdir(celeba_dir):
            # This function can fail if the Google Drive link has received many
            # recent requests. One may have to download and unzip it manually.
            download_file_from_google_drive(
                "0B7EVK8r0v71pZjFTYXZWM3FlRnM",
                celeba_dir,
                filename="celeba.zip",
            )
            extract_archive(osp.join(celeba_dir, "celeba.zip"))

            url = (
                "https://raw.githubusercontent.com/PolinaKirichenko/"
                "deep_feature_reweighting/main/celeba_metadata.csv"
            )
            wget.download(url, out=celeba_dir)

    def load_data(self):
        celeba_dir = osp.join(self.root, "celeba")
        metadata_df = pd.read_csv(osp.join(celeba_dir, "celeba_metadata.csv"))
        self.data = np.asarray(metadata_df["img_filename"].values)
        self.data = np.asarray([osp.join(celeba_dir, d) for d in self.data])

        self.targets = np.asarray(metadata_df["y"].values)
        gender = np.asarray(metadata_df["place"].values)
        nonblond = np.argwhere(self.targets == 0).flatten()
        blond = np.argwhere(self.targets == 1).flatten()
        women = np.argwhere(gender == 0).flatten()
        men = np.argwhere(gender == 1).flatten()

        self.groups = [
            np.intersect1d(nonblond, women),
            np.intersect1d(nonblond, men),
            np.intersect1d(blond, women),
            np.intersect1d(blond, men),
        ]

        split = np.asarray(metadata_df["split"].values)
        self.train_indices = np.argwhere(split == 0).flatten()
        self.val_indices = np.argwhere(split == 1).flatten()
        self.test_indices = np.argwhere(split == 2).flatten()

        # Adds group indices into targets for metrics.
        targets = []
        for j, t in enumerate(self.targets):
            g = [k for k, group in enumerate(self.groups) if j in group][0]
            targets.append([t, g])
        self.targets = np.asarray(targets)
    
class CelebA(DataModule):
    """DataModule for the CelebA dataset."""

    def __init__(self, args, **kwargs):
        super().__init__(args, CelebADataset, 2, 4, **kwargs)

    def augmented_transforms(self):
        transform = Compose([
            RandomResizedCrop((224, 224), scale=(0.7, 1.0)),
            RandomHorizontalFlip(),
            ToTensor(),
            Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ])

        return transform

    def default_transforms(self):
        transform = Compose([
            Resize((256, 256)),
            CenterCrop((224, 224)),
            ToTensor(),
            Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ])

        return transform


milkshake/datamodules/celeba.py
tmlabonte/last-layer-retraining
"""Dataset and DataModule for the CivilComments dataset."""

# Imports Python packages.
import numpy as np
import os.path as osp
import pickle
from transformers import BertTokenizer
import wilds

# Imports PyTorch packages.
import torch

# Imports milkshake packages.
from milkshake.datamodules.dataset import Dataset
from milkshake.datamodules.datamodule import DataModule
from milkshake.utils import to_np


class CivilCommentsDataset(Dataset):
    """Dataset for the CivilComments dataset."""

    def __init__(self, *xargs, **kwargs):
        super().__init__(*xargs, **kwargs)

    def download(self):
        pass

    def load_data(self):
        dataset = wilds.get_dataset(
            dataset="civilcomments",
            download=True,
            root_dir=self.root,
        )

        spurious_names = ["male", "female", "LGBTQ", "black", "white",
                          "christian", "muslim", "other_religions"]
        column_names = dataset.metadata_fields
        spurious_cols = [column_names.index(name) for name in spurious_names]
        spurious = to_np(dataset._metadata_array[:, spurious_cols].sum(-1).clip(max=1))

        prefix = osp.join(self.root, "civilcomments_v1.0")
        data_file = osp.join(prefix, "civilcomments_token_data.pt")
        targets_file = osp.join(prefix, "civilcomments_token_targets.pt")

        if not osp.isfile(data_file):
            tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
            def tokenize(text):
                tokens = tokenizer(
                    text,
                    padding="max_length",
                    truncation=True,
                    max_length=220,
                    return_tensors="pt",
                )

                return torch.squeeze(torch.stack((
                    tokens["input_ids"], tokens["attention_mask"], 
                    tokens["token_type_ids"]), dim=2), dim=0)

            data = []
            targets = []
            ln = len(dataset)
            for j, d in enumerate(dataset):
                print(f"Caching {j}/{ln}")
                data.append(tokenize(d[0]))
                targets.append(d[1])
            data = torch.stack(data)
            targets = torch.stack(targets)
            torch.save(data, data_file)
            torch.save(targets, targets_file)

        self.data = torch.load(data_file).numpy()
        self.targets = torch.load(targets_file).numpy()

        self.groups = [
            np.intersect1d((~self.targets+2).nonzero()[0], (~spurious+2).nonzero()[0]),
            np.intersect1d((~self.targets+2).nonzero()[0], spurious.nonzero()[0]),
            np.intersect1d(self.targets.nonzero()[0], (~spurious+2).nonzero()[0]),
            np.intersect1d(self.targets.nonzero()[0], spurious.nonzero()[0]),
        ]
        
        split = dataset._split_array
        self.train_indices = np.argwhere(split == 0).flatten()
        self.val_indices = np.argwhere(split == 1).flatten()
        self.test_indices = np.argwhere(split == 2).flatten()

        # Adds group indices into targets for metrics.
        targets = []
        for j, t in enumerate(self.targets):
            g = [k for k, group in enumerate(self.groups) if j in group][0]
            targets.append([t, g])
        self.targets = np.asarray(targets)

class CivilComments(DataModule):
    """DataModule for the CivilComments dataset."""

    def __init__(self, args, **kwargs):
        super().__init__(args, CivilCommentsDataset, 2, 4, **kwargs)

    def augmented_transforms(self):
        return None

    def default_transforms(self):
        return None

milkshake/datamodules/civilcomments.py
tmlabonte/last-layer-retraining
"""Parent class for a classification datamodule."""

# Imports Python builtins.
from abc import abstractmethod
from copy import deepcopy
import random

# Imports Python packages.
import numpy as np

# Imports PyTorch packages.
from torch import Generator, randperm
from torch.utils.data import DataLoader, WeightedRandomSampler
from pl_bolts.datamodules.vision_datamodule import VisionDataModule

# Imports milkshake packages.
from milkshake.datamodules.dataset import Subset
from milkshake.utils import random_split


class DataModule(VisionDataModule):
    """Parent class for a classification datamodule.

    Extends the basic PL.VisionDataModule to play nice with
    torchvision.datasets.VisionDataset and adds some custom functionality
    such as label noise, balanced sampling, fixed splits, etc.

    It inherits from PL.VisionDataModule, but milkshake.DataModule is more general
    and can be used for vision, language, or other applications.

    Attributes:
        dataset_class: A milkshake.datamodules.datasets.Dataset class.
        num_classes: The number of classes.
        num_groups: The number of groups.
        balanced_sampler: Whether to use a class-balanced random sampler during training.
        data_augmentation: Whether to use data augmentation during training.
        label_noise: Whether to add label noise during training.
        dataset_train: A milkshake.datamodules.dataset.Subset for training
        dataset_val: A milkshake.datamodules.dataset.Subset for validation.
        dataset_test: A milkshake.datamodules.dataset.Subset for testing.
        train_transforms: Preprocessing for the train set data.
        val_transforms: Preprocessing for the validation set data.
        test_transforms: Preprocessing for the test set data.
        train_transforms: Preprocessing for the train set targets.
        val_target_transforms: Preprocessing for the validation set targets.
        test_target_transforms: Preprocessing for the test set targets.
    """

    def __init__(self, args, dataset_class, num_classes, num_groups):
        """Initializes a DataModule and sets transforms.

        Args:
            args: The configuration dictionary.
            dataset_class: A milkshake.datamodules.datasets.Dataset class.
            num_classes: The number of classes.
            num_groups: The number of groups.
        """

        super().__init__(
            batch_size=args.batch_size,
            data_dir=args.data_dir,
            drop_last=False,
            normalize=True,
            num_workers=args.num_workers,
            pin_memory=True,
            seed=args.seed,
            shuffle=True,
            val_split=args.val_split,
        )

        self.persistent_workers = args.persistent_workers

        self.dataset_class = dataset_class
        self.num_classes = num_classes
        self.num_groups = num_groups
        
        self.balanced_sampler = args.balanced_sampler
        self.data_augmentation = args.data_augmentation
        self.label_noise = args.label_noise
         
        self.train_transforms = self.default_transforms()
        self.val_transforms = self.default_transforms()
        self.test_transforms = self.default_transforms()

        self.train_target_transforms = None
        self.val_target_transforms = None
        self.test_target_transforms = None

        if self.data_augmentation:
            self.train_transforms = self.augmented_transforms()

    @abstractmethod
    def augmented_transforms(self):
        """Returns torchvision.transforms for use with data augmentation."""

    @abstractmethod
    def default_transforms(self):
        """Returns default torchvision.transforms."""
 
    def prepare_data(self):
        """Downloads datasets to disk if necessary."""

        self.dataset_class(self.data_dir, train=True, download=True)
        self.dataset_class(self.data_dir, train=False, download=True)

    def load_msg(self):
        """Returns a descriptive message about the DataModule configuration."""

        msg = f"Loading {type(self).__name__}"
 
        if hasattr(self, "dataset_test"):
            msg = msg + " with test split."
        elif hasattr(self, "dataset_val"):
            if self.dataset_val.val_indices is not None:
                msg = msg + " with preset val split."
            else:
                msg = msg + f" with {int(self.val_split * 100)}% random val split."

            if self.data_augmentation:
                msg = msg[:-1] + " and data augmentation."
            if self.label_noise:
                msg = msg[:-1] + f" and {int(self.label_noise * 100)}% label noise."
            if self.balanced_sampler:
                msg = msg[:-1] + " and a balanced sampler."

        return msg

    def train_preprocess(self, dataset_train):
        """Preprocesses training dataset. Here, injects label noise if desired.

        Args:
            dataset_train: A milkshake.datamodules.dataset.Subset for training.

        Returns:
            The modified training dataset.
        """

        if self.label_noise:
            train_indices = dataset_train.train_indices
            num_labels = len(train_indices)
            num_noised_labels = int(self.label_noise * num_labels)

            for i, target in enumerate(train_indices[:num_noised_labels]):
                labels = [j for j in range(self.num_classes) if j != i]
                dataset_train.targets[i] = random.choice(labels)

        return dataset_train

    def val_preprocess(self, dataset_val):
        """Preprocesses validation dataset. Does nothing here, but can be overriden.

        Args:
            dataset_val: A milkshake.datamodules.dataset.Subset for validation.

        Returns:
            The modified validation dataset.
        """

        return dataset_val

    def test_preprocess(self, dataset_test):
        """Preprocesses test dataset. Does nothing here, but can be overidden.

        Args:
            dataset_test: A milkshake.datamodules.dataset.Subset for testing.

        Returns:
            The modified test dataset.
        """

        return dataset_test

    def setup(self, stage=None):
        """Instantiates and preprocesses datasets.
        
        Args:
            stage: "train", "validate", or "test".
        """

        if stage == "test":
            dataset_test = self.dataset_class(
                self.data_dir,
                train=False,
                transform=self.test_transforms,
                target_transform=self.test_target_transforms,
            )

            dataset_test = self.test_preprocess(dataset_test)
            self.dataset_test = dataset_test
        else:
            dataset_train = self.dataset_class(
                self.data_dir,
                train=True,
                transform=self.train_transforms,
                target_transform=self.train_target_transforms,
            )

            dataset_val = self.dataset_class(
                self.data_dir,
                train=True,
                transform=self.val_transforms,
                target_transform=self.val_target_transforms,
            )

            
            dataset_train = self._split_dataset(dataset_train)
            self.dataset_train = self.train_preprocess(dataset_train)

            dataset_val = self._split_dataset(dataset_val, val=True)
            self.dataset_val = self.val_preprocess(dataset_val)

        print(self.load_msg())

    def _split_dataset(self, dataset, val=False):
        """Splits dataset into training and validation subsets.

        Checks for a preset split, then returns a random split if it does not exist.
        
        Args:
            dataset: A milkshake.datamodules.dataset.Dataset.
            val: Whether to return the validation set (otherwise returns the training set).

        Returns:
            A milkshake.datamodules.dataset.Subset of the given dataset with the desired split.
        """

        if dataset.train_indices is None or dataset.val_indices is None:
            len_dataset = len(dataset)
            splits = self._get_splits(len_dataset)

            # Computes a random split with proportion based on self.split.
            dataset_train, dataset_val = random_split(
                dataset,
                splits,
                generator=Generator().manual_seed(self.seed),
            )
        else:
            dataset_train = Subset(dataset, dataset.train_indices)
            dataset_val = Subset(dataset, dataset.val_indices)
                        
        if val:
            return dataset_val
        return dataset_train

    def train_dataloader(self):
        """Returns DataLoader for the train dataset."""

        # Instantiates class-balanced random sampler.
        if self.balanced_sampler:
            indices = self.dataset_train.train_indices
            targets = self.dataset_train.targets[indices]

            # Removes group dimension if necessary.
            if targets[0].ndim > 0:
                targets = targets[:, 0]

            counts = np.bincount(targets)
            label_weights = 1. / counts
            weights = label_weights[targets]
            sampler = WeightedRandomSampler(weights, len(weights))

            return self._data_loader(self.dataset_train, sampler=sampler)

        return self._data_loader(self.dataset_train, shuffle=self.shuffle)

    def val_dataloader(self):
        """Returns DataLoader(s) for the val dataset."""

        return self._data_loader(self.dataset_val)

    def test_dataloader(self):
        """Returns DataLoader(s) for the test dataset."""

        return self._data_loader(self.dataset_test)

    def _data_loader(self, dataset, shuffle=False, sampler=None):
        """Instantiates DataLoader with the given dataset.

        Args:
            dataset: A milkshake.datamodules.dataset.Dataset.
            shuffle: Whether to shuffle the data indices.
            sampler: A torch.utils.data.Sampler for selecting data indices.

        Returns:
            A torch.utils.data.DataLoader with the given configuration.
        """

        return DataLoader(
            dataset,
            batch_size=self.batch_size,
            drop_last=self.drop_last,
            num_workers=self.num_workers,
            pin_memory=self.pin_memory,
            persistent_workers=self.persistent_workers,
            sampler=sampler,
            shuffle=shuffle,
        )


milkshake/datamodules/datamodule.py
tmlabonte/last-layer-retraining
"""Parent class for a classification dataset."""

# Imports Python builtins
from abc import abstractmethod

# Imports Python packages.
import numpy as np
import os.path as osp
from PIL import Image

# Imports PyTorch packages.
import torch
from torchvision.datasets.vision import VisionDataset


class Dataset(VisionDataset):
    """Parent class for a classification dataset.

    Mostly the same as torchvision.datasets.vision.VisionDataset with some
    extra pieces from torchvision.datasets.vision.CIFAR10Dataset (e.g., the
    train flag) and support for multiple groups.

    It inherits from torchvision.datasets.vision.VisionDataset but milkshake.Dataset
    is more general and can be used for vision, language, or other applications.

    Attributes:
        root: The location of the dataset on disk.
        transform: Preprocessing for the data.
        target_transform: Preprocessing for the targets.
        data: np.ndarray containing np.ndarrays or string filenames.
        targets: np.ndarray containing classification targets.
        train: Whether the dataset should be loaded in train mode (for transforms, etc.).
        train_indices: Optional np.ndarray of indices of train set.
        val_indices: Optional np.ndarray of indices of val set.
        test_indices: Optional np.ndarray of indices of test set.
        group: If the dataset has multiple groups, specifies which one to initialize.
        groups: If the dataset has multiple groups, lists indices belonging to each group.
    """

    def __init__(
        self,
        root,
        train=True,
        transform=None,
        target_transform=None, 
        download=False,
    ):
        """Initializes a Dataset and downloads to disk if requested.

        If train_indices, etc. are None, then DataModule will do a random split
        based on args.val_split. The np.ndarrays self.data and self.targets
        should hold both training and validation data when train=True and only
        test data when train=False. The split is calculated in the DataModule.

        Args:
            root: The location of the dataset on disk.
            train: Whether the dataset should be loaded in train mode (for transforms, etc.).
            transform: Preprocessing for the data.
            target_transform: Preprocessing for the targets.
            download: Whether to download the dataset to disk.
        """

        VisionDataset.__init__(
            self,
            root,
            transform=transform,
            target_transform=target_transform,
        )

        self.train = train

        self.data = None
        self.targets = None
        self.train_indices = None
        self.val_indices = None
        self.test_indices = None
        self.groups = None

        if download:
            self.download()

        self.load_data()

        # Initializes Dataset with test data only.
        if (self.data is not None and self.targets is not None and
                self.test_indices is not None and not self.train):
            self.data = self.data[self.test_indices]
            self.targets = self.targets[self.test_indices]
            if self.groups is not None:
                self.groups = [
                    np.in1d(self.test_indices, group).nonzero()[0]
                    for group in self.groups
                ]
            self.train_indices = None
            self.val_indices = None
            self.test_indices = np.arange(len(self.data))

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        datum, target = self.data[index], self.targets[index]

        # Loads a PIL.Image (for vision) or a string (for NLP) from the cached
        # np.ndarray data or from disk. If a string is passed, tries to check if
        # it is a filepath; if so loads an image, otherwise returns the string.
        # TODO: Make lazy loading option for each dataset.
        if isinstance(datum, np.ndarray):
            if datum.dtype not in (str, np.str_):
                try:
                    datum = Image.fromarray(datum)
                except:
                    pass # e.g., BERT tokens in an np.ndarray
        elif isinstance(datum, (str, np.str_)):
            if osp.isfile(datum):
                datum = Image.open(datum)
        elif isinstance(datum, torch.Tensor):
            pass
        else:
            raise ValueError("Unrecognized data type.")

        if self.transform is not None:
            datum = self.transform(datum)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return datum, target

    @abstractmethod
    def download(self):
        """Downloads dataset to disk."""

    @abstractmethod
    def load_data(self):
        """Initializes self.data and self.targets, and optionally indices and groups."""

class Subset(Dataset):
    """Subset of a Dataset at specified indices.
    
    Modified from torch.utils.Subset to allow interactions as if it was a
    milkshake.datamodules.dataset.Dataset (e.g., by indices, groups, etc.).
    """

    def __init__(self, dataset, indices):
        """Initializes a Subset and sets the new indices.
        
        Args:
            dataset: A milkshake.datamodules.Dataset.
            indices: The indices to utilize in the subset.
        """

        self.indices = indices
        self.root = dataset.root
        self.transform = dataset.transform
        self.target_transform = dataset.target_transform
        self.train = dataset.train

        self.groups = dataset.groups

        self.data = dataset.data[indices]
        self.targets = dataset.targets[indices]
        
        self.train_indices = dataset.train_indices
        self.val_indices = dataset.val_indices
        self.test_indices = dataset.test_indices
        
        # Gets subsets of train_indices, etc. that are present in indices and
        # converts them to new indices taking values in [0, len(indices)].
        # If self.train_indices is None, then we are doing a random split
        # so we can set all indices to [0, len(indices)].
        if self.train_indices is not None:
            self.train_indices = np.in1d(indices, self.train_indices).nonzero()[0]
        else:
            self.train_indices = np.arange(len(indices))
            
        if self.val_indices is not None:
            self.val_indices = np.in1d(indices, self.val_indices).nonzero()[0]
        else:
            self.val_indices = np.arange(len(indices))
            
        if not self.train and self.test_indices is not None:
            self.test_indices = np.in1d(indices, self.test_indices).nonzero()[0]
        elif not self.train:
            self.test_indices = np.arange(len(indices))
            
        if self.groups is not None:
            self.groups = [
                np.in1d(indices, group).nonzero()[0]
                for group in dataset.groups
            ]

    def __getitem__(self, index):
        return super().__getitem__(index)

    def __len__(self):
        return len(self.indices)


milkshake/datamodules/dataset.py
tmlabonte/last-layer-retraining
"""DataModule for disagreement-based datasets."""

# Imports Python packages.
import numpy as np
from numpy.random import default_rng

# Imports PyTorch packages.
import torch
import torch.nn.functional as F
from torch.utils.data import WeightedRandomSampler

# Imports milkshake packages.
from milkshake.datamodules.dataset import Subset
from milkshake.datamodules.datamodule import DataModule
from milkshake.utils import to_np


class Disagreement(DataModule):
    """DataModule for disagreement sets used for last-layer retraining.

    The original DFR procedure uses group annotations to construct a reweighting
    dataset that has equal data from every group. We propose using disagreement
    between the ERM model and a regularized (e.g., dropout or early-stopped) 
    model as an alternative. This enables construction of nearly-group-balanced
    reweighting datasets without the need for group annotations.

    This class is currently only defined for datasets with a preset val split,
    i.e., milkshake.datasets which have train_indices and val_indices.

    Attributes:
    """

    def __init__(
        self,
        args,
        *xargs,
        early_stop_model=None,
        model=None,
        num_data=None,
        worst_group_pct=None,
    ):
        """Initializes a Disagreement DataModule.
        
        Args:
            args: The configuration dictionary.
            *xargs: Any additional positional arguments.
            model: The milkshake.models.Model used for disagreement.
        """

        super().__init__(args, *xargs)
        
        self.heldout_pct = 0.5

        self.balance_finetune = args.balance_finetune
        self.finetune_type = args.finetune_type if hasattr(args, "finetune_type") else None
        self.split = args.split
        self.train_pct = args.train_pct
        self.val_pct = args.val_pct

        self.early_stop_model = early_stop_model.cuda() if early_stop_model else None
        self.model = model.cuda() if model else None
        self.num_data = num_data
        self.worst_group_pct = worst_group_pct

    def _split_dataset(self, dataset_train, dataset_val):
        """Splits dataset into training and validation subsets.

        Args:
            dataset: A milkshake.datamodules.Dataset.

        Returns:
        """

        if dataset_train.train_indices is not None:
            train_inds = dataset_train.train_indices
            val_inds = dataset_val.val_indices
            default_rng(seed=self.seed).shuffle(val_inds)
        else:
            len_dataset = len(dataset_train)
            len_train, _ = self._get_splits(len_dataset)
            inds = np.arange(len_dataset)
            default_rng(seed=self.seed).shuffle(inds)
            train_inds = inds[:len_train]
            val_inds = inds[len_train:]
            dataset_train.train_indices = train_inds
            dataset_val.train_indices = train_inds
            dataset_train.val_indices = val_inds
            dataset_val.val_indices = val_inds

        disagreement_num = int(self.heldout_pct * len(val_inds))
        new_val_inds = val_inds[disagreement_num:]

        if self.split == "train" and self.train_pct == 100:
            new_train_inds = train_inds
            new_dis_inds = val_inds[:disagreement_num]
        elif self.split == "train":
            default_rng(seed=self.seed).shuffle(train_inds)
            train_num = int(len(train_inds) * self.train_pct / 100)
            new_train_inds = train_inds[:train_num]
            new_dis_inds = train_inds[train_num:]
        elif self.split == "combined":
            combined_inds = np.concatenate((train_inds, val_inds))
            train_num = int((len(train_inds) + disagreement_num) * self.train_pct / 100)
            new_dis_num = len(combined_inds) - int((1 - self.heldout_pct) * len(val_inds))
            new_combined_inds = combined_inds[:new_dis_num]

            default_rng(seed=self.seed).shuffle(new_combined_inds)
            new_train_inds = new_combined_inds[:train_num]
            new_dis_inds = new_combined_inds[train_num:]

        dataset_t = Subset(dataset_train, new_train_inds)
        dataset_d = Subset(dataset_val, new_dis_inds)
        dataset_v = Subset(dataset_val, new_val_inds)

        if self.split == "combined":
            dataset_t.train_indices = list(range(len(dataset_t)))

        dataset_d.val_indices = new_dis_inds
        dataset_v.val_indices = new_val_inds

        return dataset_t, dataset_d, dataset_v

    def train_dataloader(self):
        """Returns DataLoader for the train dataset (after disagreement).
        
           This method also handles class- and group-balancing, either by
           sampling the data to ensure balanced minibatches ("sampler")
           or by training on a balanced subset of the data ("subset").
        """
        
        if self.finetune_type == "group-unbalanced retraining":
            if self.worst_group_pct:
                # Sets the worst groups (zero-indexed).
                min_group_nums = {
                    "Waterbirds": np.array([1, 2]),
                    "CelebA": np.array([3]),
                    "CivilComments": np.array([3]),
                    "MultiNLI": np.array([3, 5]),
                }

                maj_group_nums = {
                    "Waterbirds": np.array([0, 3]),
                    "CelebA": np.array([2]),
                    "CivilComments": np.array([2]),
                    "MultiNLI": np.array([2, 4]),
                }

                dataset = self.__class__.__bases__[0].__name__[:-12]
                min_groups = min_group_nums[dataset]
                maj_groups = maj_group_nums[dataset]
                totals = np.array([len(x) for x in self.dataset_train.groups])
                num_groups = len(totals)
                num_maj_groups = len(maj_groups)

                smallest_min_group = min([t for j, t in enumerate(totals)
                                          if j in min_groups])
                smallest_maj_group = min([t for j, t in enumerate(totals)
                                          if j in maj_groups])

                downsample_num = min(smallest_maj_group // 2, smallest_min_group)
                totals = [downsample_num] * num_groups

                """
                # Makes CelebA and CivilComments class-unbalanced.
                maj_ratio = sum(totals[:2]) / sum(totals)
                maj = int(-(maj_ratio * downsample_num) / (maj_ratio - 1))
                totals = [maj, maj, downsample_num, downsample_num]
                """

                removed = 0
                for j, _ in enumerate(totals):
                    if j in min_groups:
                        new = int(totals[j] * self.worst_group_pct / 100)
                        removed += totals[j] - new
                        totals[j] = new

                maj_per_group = removed // num_maj_groups
                for j, t in enumerate(totals):
                    if j in maj_groups:
                        totals[j] += maj_per_group

                indices = self.dataset_train.train_indices
                defaut_rng(seed=self.seed).shuffle(indices)

                new_indices = []
                nums = [0] * num_groups
                for x in indices:
                    for j, group in enumerate(self.dataset_train.groups):
                        if x in group and nums[j] < totals[j]:
                            new_indices.append(x) 
                            nums[j] += 1
                            break
                print(nums)

                new_indices = np.array(new_indices)
                self.dataset_train = Subset(self.dataset_train, new_indices)

                self.balanced_sampler = False
                return super().train_dataloader()
            elif self.balance_finetune == "sampler":
                self.balanced_sampler = True
                return super().train_dataloader()
            elif self.balance_finetune == "subset":
                indices = self.dataset_train.train_indices
                targets = self.dataset_train.targets[indices]

                # Shuffles indices and targets in unison.
                p = default_rng(seed=self.seed).permutation(len(indices))
                indices = indices[p]
                targets = targets[p]

                min_target = min(np.unique(targets, return_counts=True)[1])
                counts = [0] * self.num_classes

                subset = []
                for idx, target in zip(indices, targets):
                    if counts[target] < min_target:
                        subset.append(idx)
                        counts[target] += 1
                print(f"Balanced subset: {counts}")

                self.dataset_train = Subset(self.dataset_train, subset)

                self.balanced_sampler=False
                return super().train_dataloader()
            elif self.balance_finetune == "none":
                self.balanced_sampler = False
                return super().train_dataloader()
        elif self.finetune_type == "group-balanced retraining":
            indices = self.dataset_train.train_indices
            groups = np.zeros(len(indices), dtype=np.int32)
            for i, x in enumerate(indices):
                for j, group in enumerate(self.dataset_train.groups):
                    if x in group:
                        groups[i] = j

            if self.balance_finetune == "sampler":
                counts = np.bincount(groups)
                label_weights = 1. / counts
                weights = label_weights[groups]
                sampler = WeightedRandomSampler(weights, len(weights))
                return self._data_loader(self.dataset_train, sampler=sampler)
            elif self.balance_finetune == "subset":
                # Shuffles indices and groups in unison.
                p = default_rng(seed=self.seed).permutation(len(indices))
                indices = indices[p]
                groups = groups[p]

                min_group = min([len(x) for x in self.dataset_train.groups])
                counts = [0] * len(self.dataset_train.groups)

                subset = []
                for idx, group in zip(indices, groups):
                    if counts[group] < min_group:
                        subset.append(idx)
                        counts[group] += 1
                print(f"Balanced subset: {counts}")

                self.dataset_train = Subset(self.dataset_train, subset)

                self.balanced_sampler = False
                return super().train_dataloader()
            elif self.balance_finetune == "none":
                raise ValueError("Set balance_finetune for group-balanced retraining.")
        else:
            if self.balance_finetune == "sampler":
                self.balanced_sampler = True
                return super().train_dataloader()
            elif self.balance_finetune == "subset":
                indices = self.dataset_train.train_indices
                targets = self.dataset_train.targets[indices]

                # Shuffles indices and targets in unison.
                p = default_rng(seed=self.seed).permutation(len(indices))
                indices = indices[p]
                targets = targets[p]

                min_target = min(np.unique(targets, return_counts=True)[1])
                counts = [0] * self.num_classes

                subset = []
                for idx, target in zip(indices, targets):
                    if counts[target] < min_target:
                        subset.append(idx)
                        counts[target] += 1

                self.dataset_train = Subset(self.dataset_train, subset)

                self.balanced_sampler=False
                return super().train_dataloader()
            elif self.balance_finetune == "none":
                self.balanced_sampler = False
                return super().train_dataloader()

    def val_dataloader(self):
        # Randomly samples specified percent of validation set (for ablations).
        indices = np.arange(len(self.dataset_val.val_indices))
        self.dataset_val.val_indices = indices
        p = default_rng(seed=self.seed).permutation(len(indices))
        total = int(self.val_pct / 100. * len(indices))
        self.dataset_val = Subset(self.dataset_val, indices[p][:total])

        dataloaders = super().val_dataloader()
        return dataloaders

    def test_dataloader(self):
        dataloaders = super().test_dataloader()
        return dataloaders

    def disagreement_dataloader(self):
        """Returns DataLoader for the disagreement set."""

        return self._data_loader(self.dataset_disagreement)

    def print_disagreements_by_group(self, dataset, all_inds, disagree=None):
        """Prints number of disagreements occuring in each group.
        
        Args:
            dataset: A milkshake.datamodules.Dataset.
            all_inds: An np.ndarray of all indices in the disagreement set.
            disagree: An optional np.ndarray of all disagreed indices.
        """

        labels_and_inds = zip(
            ("All", "Disagreements"),
            (all_inds, disagree),
        )

        print("Disagreements by group")
        for label, inds in labels_and_inds:
            if inds is not None:
                nums = []
                for group in dataset.groups:
                    nums.append(len(np.intersect1d(inds, group)))
                print(f"{label}: {nums}")
    
    def disagreement(self):
        """Computes finetuning set and saves it as self.dataset_train.
        
        self.dataset_disagreement is initially the self.heldout_pct of the
        held-out set. Here, we perform some computation (i.e., the actual
        disagreement) on self.dataset_disagreement to get indices for
        finetuning. Then, we set these indices as self.dataset_train.
        """

        dataloader = self.disagreement_dataloader()
        batch_size = dataloader.batch_size

        all_inds = dataloader.dataset.val_indices

        new_set = self.dataset_class(
            self.data_dir,
            train=True,
            transform=self.train_transforms,
        )

        if new_set.train_indices is not None and new_set.val_indices is not None:
            new_set.train_indices = np.concatenate((new_set.train_indices, new_set.val_indices))
        else:
            new_set.train_indices = np.arange(len(new_set))

        if "retraining" in self.finetune_type:
            self.dataset_train = Subset(new_set, all_inds)
            return

        all_orig_logits = []
        all_logits = []
        all_orig_probs = []
        all_probs = []
        all_targets = []

        # Performs misclassification or disagreement with self.model.
        with torch.no_grad():
            for i, batch in enumerate(dataloader):
                inputs, targets = batch
                inputs = inputs.cuda()
                targets = targets.cuda()

                if "early-stop" in self.finetune_type:
                    self.early_stop_model.eval()
                    orig_logits = self.early_stop_model(inputs)
                    orig_probs = F.softmax(orig_logits, dim=1)
                    orig_preds = torch.argmax(orig_probs, dim=1)
                else:
                    self.model.eval()
                    orig_logits = self.model(inputs)
                    orig_probs = F.softmax(orig_logits, dim=1)
                    orig_preds = torch.argmax(orig_probs, dim=1)

                if "misclassification" in self.finetune_type:
                    logits = self.model(inputs)
                else:
                    if "dropout" in self.finetune_type:
                        self.model.train()
                        logits = self.model(inputs)
                    elif "early-stop" in self.finetune_type:
                        logits = self.early_stop_model(inputs)

                probs = F.softmax(logits, dim=1)
                preds = torch.argmax(probs, dim=1)

                all_orig_logits.append(orig_logits)
                all_logits.append(logits)
                all_orig_probs.append(orig_probs)
                all_probs.append(probs)
                all_targets.append(targets)

        all_orig_logits = torch.cat(all_orig_logits)
        all_logits = torch.cat(all_logits)
        all_orig_probs = torch.cat(all_orig_probs)
        all_probs = torch.cat(all_probs)
        all_targets = torch.cat(all_targets)

        # Removes group dimension.
        if all_targets[0].ndim > 0:
            all_targets = all_targets[:, 0]
        else:
            all_targets = all_targets
        
        kldiv = F.kl_div(torch.log(all_probs), all_orig_probs, reduction="none")
        kldiv = torch.mean(kldiv, dim=1).squeeze()
        loss = F.cross_entropy(all_orig_logits, all_targets, reduction="none").squeeze()

        if "misclassification" in self.finetune_type:
            disagreements = to_np(torch.topk(loss, k=self.num_data)[1])
        elif "disagreement" in self.finetune_type:
            disagreements = to_np(torch.topk(kldiv, k=self.num_data)[1])

        disagree = all_inds[disagreements]
            
        self.dataset_train = Subset(new_set, disagree)
        self.print_disagreements_by_group(new_set, all_inds, disagree=disagree)
        
    def setup(self, stage=None):
        """Instantiates and preprocesses datasets.

        Performs disagreement if self.model (i.e., the model which calculates
        disagreements) is specified.
        """

        dataset_train = self.dataset_class(
            self.data_dir,
            train=True,
            transform=self.train_transforms,
        )

        dataset_val = self.dataset_class(
            self.data_dir,
            train=True,
            transform=self.val_transforms,
        )

        dataset_test = self.dataset_class(
            self.data_dir,
            train=False,
            transform=self.test_transforms,
        )

        # Creates disagreement sets in addition to regular train/val split.
        dataset_train = self.train_preprocess(dataset_train)
        dataset_val = self.val_preprocess(dataset_val)
        self.dataset_train, self.dataset_disagreement, self.dataset_val = \
            self._split_dataset(dataset_train, dataset_val)

        dataset_test = self.test_preprocess(dataset_test)
        self.dataset_test = dataset_test
        
        if stage == "fit":
            # Performs disagreement and sets new train dataset.
            if self.model:
                print("Computing disagreements...")
                self.disagreement()
                del self.model


milkshake/datamodules/disagreement.py
tmlabonte/last-layer-retraining
"""Dataset and DataModule for the MultiNLI dataset."""

# Imports Python builtins.
import os
import os.path as osp
import sys

# Imports Python packages.
import numpy as np
import pandas as pd
import wget

# Imports PyTorch packages.
import torch
from torchvision.datasets.utils import (
    extract_archive,
)

# Imports milkshake packages.
from milkshake.datamodules.dataset import Dataset
from milkshake.datamodules.datamodule import DataModule


class MultiNLIDataset(Dataset):
    """Dataset for the MultiNLI dataset."""

    def __init__(self, *xargs, **kwargs):
        super().__init__(*xargs, **kwargs)

    def download(self):
        multinli_dir = osp.join(self.root, "multinli")
        if not osp.isdir(multinli_dir):
            os.makedirs(multinli_dir)

            url = (
                "https://github.com/kohpangwei/group_DRO/raw/"
                "f7eae929bf4f9b3c381fae6b1b53ab4c6c911a0e/"
                "dataset_metadata/multinli/metadata_random.csv"
            )
            wget.download(url, out=multinli_dir)

            url = "https://nlp.stanford.edu/data/dro/multinli_bert_features.tar.gz"
            wget.download(url, out=multinli_dir)
            extract_archive(osp.join(multinli_dir, "multinli_bert_features.tar.gz"))

            url = (
                "https://raw.githubusercontent.com/izmailovpavel/"
                "spurious_feature_learning/6d098440c697a1175de6a24"
                "d7a46ddf91786804c/dataset_files/utils_glue.py"
            )
            wget.download(url, out=multinli_dir)
    
    def load_data(self):
        multinli_dir = osp.join(self.root, "multinli")
        sys.path.append(multinli_dir)
        metadata_path = osp.join(multinli_dir, "metadata_random.csv")
        metadata_df = pd.read_csv(metadata_path)

        bert_filenames = [
            "cached_train_bert-base-uncased_128_mnli",
            "cached_dev_bert-base-uncased_128_mnli",
            "cached_dev_bert-base-uncased_128_mnli-mm",
        ]

        features_array = sum([torch.load(osp.join(multinli_dir, name))
                              for name in bert_filenames], start=[])

        all_input_ids = torch.tensor([
            f.input_ids for f in features_array
        ]).long()
        all_input_masks = torch.tensor([
            f.input_mask for f in features_array
        ]).long()
        all_segment_ids = torch.tensor([
            f.segment_ids for f in features_array
        ]).long()

        self.data = torch.stack((
            all_input_ids,
            all_input_masks,
            all_segment_ids,
        ), dim=2)
        self.targets = np.asarray(metadata_df["gold_label"].values)

        spurious = np.asarray(metadata_df["sentence2_has_negation"].values)
        no_negation = np.argwhere(spurious == 0).flatten()
        negation = np.argwhere(spurious == 1).flatten()
        contradiction = np.argwhere(self.targets == 0).flatten()
        entailment = np.argwhere(self.targets == 1).flatten()
        neutral = np.argwhere(self.targets == 2).flatten()

        self.groups = [
            np.intersect1d(contradiction, no_negation),
            np.intersect1d(contradiction, negation),
            np.intersect1d(entailment, no_negation),
            np.intersect1d(entailment, negation),
            np.intersect1d(neutral, no_negation),
            np.intersect1d(neutral, negation),
        ]

        split = np.asarray(metadata_df["split"].values)
        self.train_indices = np.argwhere(split == 0).flatten()
        self.val_indices = np.argwhere(split == 1).flatten()
        self.test_indices = np.argwhere(split == 2).flatten()

        # Adds group indices into targets for metrics.
        targets = []
        for j, t in enumerate(self.targets):
            g = [k for k, group in enumerate(self.groups) if j in group][0]
            targets.append([t, g])
        self.targets = np.asarray(targets)

class MultiNLI(DataModule):
    """DataModule for the MultiNLI dataset."""

    def __init__(self, args, **kwargs):
        super().__init__(args, MultiNLIDataset, 3, 6, **kwargs)

    def augmented_transforms(self):
        return None

    def default_transforms(self):
        return None

milkshake/datamodules/multinli.py
tmlabonte/last-layer-retraining
"""Dataset and DataModule for the Waterbirds dataset."""

# Imports Python builtins.
import os.path as osp

# Imports Python packages.
import numpy as np
import pandas as pd
from PIL import Image

# Imports PyTorch packages.
import torch
import torch.nn.functional as F
from torchvision.datasets.utils import download_and_extract_archive
from torchvision.transforms import (
    CenterCrop,
    Compose,
    Normalize,
    RandomHorizontalFlip,
    RandomResizedCrop,
    Resize,
    ToTensor,
)

# Imports milkshake packages.
from milkshake.datamodules.dataset import Dataset
from milkshake.datamodules.datamodule import DataModule


class WaterbirdsDataset(Dataset):
    """Dataset for the Waterbirds dataset."""

    def __init__(self, *xargs, **kwargs):
        super().__init__(*xargs, **kwargs)

    def download(self):
        waterbirds_dir = osp.join(self.root, "waterbirds")
        if not osp.isdir(waterbirds_dir):
            url = (
                "http://worksheets.codalab.org/rest/bundles/"
                "0x505056d5cdea4e4eaa0e242cbfe2daa4/contents/blob/"
            )

            download_and_extract_archive(
                url,
                waterbirds_dir,
                filename="waterbirds.tar.gz",
            )

    def load_data(self):
        waterbirds_dir = osp.join(self.root, "waterbirds")
        metadata_df = pd.read_csv(osp.join(waterbirds_dir, "metadata.csv"))
        self.data = np.asarray(metadata_df["img_filename"].values)
        self.data = np.asarray([osp.join(waterbirds_dir, d) for d in self.data])

        self.targets = np.asarray(metadata_df["y"].values)
        background = np.asarray(metadata_df["place"].values)
        landbirds = np.argwhere(self.targets == 0).flatten()
        waterbirds = np.argwhere(self.targets == 1).flatten()
        land = np.argwhere(background == 0).flatten()
        water = np.argwhere(background == 1).flatten()

        self.groups = [
            np.intersect1d(landbirds, land),
            np.intersect1d(landbirds, water),
            np.intersect1d(waterbirds, land),
            np.intersect1d(waterbirds, water),
        ]
        
        split = np.asarray(metadata_df["split"].values)
        self.train_indices = np.argwhere(split == 0).flatten()
        self.val_indices = np.argwhere(split == 1).flatten()
        self.test_indices = np.argwhere(split == 2).flatten()

        # Adds group indices into targets for metrics.
        targets = []
        for j, t in enumerate(self.targets):
            g = [k for k, group in enumerate(self.groups) if j in group][0]
            targets.append([t, g])
        self.targets = np.asarray(targets)

class Waterbirds(DataModule):
    """DataModule for the Waterbirds dataset."""

    def __init__(self, args, **kwargs):
        super().__init__(args, WaterbirdsDataset, 2, 4, **kwargs)

    def augmented_transforms(self):
        transform = Compose([
            RandomResizedCrop((224, 224), scale=(0.7, 1.0)),
            RandomHorizontalFlip(),
            ToTensor(),
            Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ])

        return transform

    def default_transforms(self):
        transform = Compose([
            Resize((256, 256)),
            CenterCrop((224, 224)),
            ToTensor(),
            Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ])

        return transform


milkshake/datamodules/waterbirds.py
tmlabonte/last-layer-retraining
"""Import functions for milkshake."""

# Imports milkshake packages.
import milkshake
from milkshake.datamodules import *
from milkshake.models import *


def valid_model_and_datamodule_names():
    """Returns valid input names for models and datamodules."""

    model_names = [n for n in milkshake.models.__all__ if n != "model"]
    datamodule_names = [n for n in milkshake.datamodules.__all__ if n not in ("dataset", "datamodule")]

    return model_names, datamodule_names

def valid_models_and_datamodules():
    """Returns {name: class} dict for valid models and datamodules."""

    model_names, datamodule_names = valid_model_and_datamodule_names()

    models = [milkshake.models.__dict__[name].__dict__ for name in model_names]
    models = [dict((k.lower(), v) for k, v in d.items()) for d in models]
    models = {name: models[j][name.replace("_", "")] for j, name in enumerate(model_names)} 

    datamodules = [milkshake.datamodules.__dict__[name].__dict__ for name in datamodule_names]
    datamodules = [dict((k.lower(), v) for k, v in d.items()) for d in datamodules]
    datamodules = {name: datamodules[j][name.replace("_", "")] for j, name in enumerate(datamodule_names)} 

    return models, datamodules


milkshake/imports.py
tmlabonte/last-layer-retraining
"""Main script for training, validation, and testing."""

# Ignores nuisance warnings. Must be called first.
from milkshake.utils import ignore_warnings
ignore_warnings()

# Imports Python builtins.
from inspect import isclass
import os
import os.path as osp
import resource

# Imports Python packages.
import wandb
from PIL import ImageFile

# Imports PyTorch packages.
from lightning_lite.utilities.seed import seed_everything
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import ModelCheckpoint, RichProgressBar
from pytorch_lightning.loggers import WandbLogger
import torch

# Imports milkshake packages.
from milkshake.args import parse_args
from milkshake.imports import valid_models_and_datamodules

# Sets matmul precision to take advantage of Tensor Cores.
torch.set_float32_matmul_precision("high")

# Prevents PIL from throwing invalid error on large image files.
ImageFile.LOAD_TRUNCATED_IMAGES = True

# Prevents DataLoader memory error.
rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (8192, rlimit[1]))

# Silences wandb printouts.
os.environ["WANDB_SILENT"] = "true"


def load_weights(args, model):
    """Loads model weights or resumes training checkpoint.

    Args:
        args: The configuration dictionary.
        model: A model which inherits from milkshake.models.Model.

    Returns:
        The input model, possibly with the state dict loaded.
    """

    args.ckpt_path = None
    if args.weights:
        if args.resume_training:
            # Resumes training state (weights, optimizer, epoch, etc.) from args.weights.
            args.ckpt_path = args.weights
            print(f"Resuming training state from {args.weights}.")
        else:
            # Loads just the weights from args.weights.
            checkpoint = torch.load(args.weights, map_location="cpu")
            state_dict = checkpoint["state_dict"]
            model.load_state_dict(state_dict, strict=False)
            print(f"Weights loaded from {args.weights}.")     

    return model

def load_trainer(args, addtl_callbacks=None):
    """Loads PL Trainer for training and validation.

    Args:
        args: The configuration dictionary.
        addtl_callbacks: Desired callbacks besides ModelCheckpoint and TQDMProgressBar.

    Returns:
        An instance of pytorch_lightning.Trainer parameterized by args.
        
    Raises:
        ValueError: addtl_callbacks is not None and not a list.
    """

    # Checking validation by epochs and by steps are mutually exclusive.
    if args.val_check_interval:
        args.check_val_every_n_epoch = None

    # Checkpointing by epochs and by steps are mutually exclusive.
    if args.ckpt_every_n_steps:
        args.ckpt_every_n_epochs = None

    # Checkpoints model at the specified number of epochs.
    checkpointer1 = ModelCheckpoint(
        filename="{epoch:02d}-{val_loss:.3f}-{val_acc:.3f}",
        save_top_k=-1,
        every_n_epochs=args.ckpt_every_n_epochs,
        every_n_train_steps=args.ckpt_every_n_steps,
    )

    # Checkpoints model with respect to validation loss.
    checkpointer2 = ModelCheckpoint(
        filename="best-{epoch:02d}-{val_loss:.3f}-{val_acc:.3f}",
        monitor="val_loss",
        every_n_epochs=args.ckpt_every_n_epochs,
        every_n_train_steps=args.ckpt_every_n_steps,
    )

    progress_bar = RichProgressBar(refresh_rate=args.refresh_rate)
 
    callbacks = [checkpointer1, checkpointer2, progress_bar]
    if addtl_callbacks is not None:
        if not isinstance(addtl_callbacks, list):
            raise ValueError("addtl_callbacks should be None or a list.")
        callbacks.extend(addtl_callbacks)

    logger = True # Activates TensorBoardLogger by default.
    if args.wandb:
        os.makedirs(args.wandb_dir, exist_ok=True)
        logger = WandbLogger(save_dir=args.wandb_dir, log_model="all")

    # Sets DDP strategy for multi-GPU training.
    # Important note: running multiple fit/validate/test loops using ddp (like
    # the Milkshake example experiments) is an anti-pattern in PyTorch Lightning
    # (see https://github.com/Lightning-AI/lightning/issues/8375). However, ddp
    # is much faster than the alternative, ddp_spawn. Therefore, we recommend
    # using ddp for single training jobs and ddp_spawn for experiments.
    args.strategy = args.strategy if args.devices > 1 else None

    trainer = Trainer.from_argparse_args(
        args,
        callbacks=callbacks,
        deterministic=True,
        logger=logger,
    )

    return trainer

def main(
    args,
    model_or_model_class,
    datamodule_or_datamodule_class,
    callbacks=None,
    model_hooks=None,
    verbose=True,
):
    """Main method for training and validation.

    Args:
        args: The configuration dictionary.
        model_or_model_class: A model or class which inherits from milkshake.models.Model.
        datamodule_or_datamodule_class: A datamodule or class which inherits from milkshake.datamodules.DataModule.
        callbacks: Any desired callbacks besides ModelCheckpoint and TQDMProgressBar.
        model_hooks: Any desired functions to run on the model before training.
        verbose: Whether to print the validation and test metrics.

    Returns:
        The trained model with its validation and test metrics.
    """

    os.makedirs(args.out_dir, exist_ok=True)
    args.devices = int(args.devices)

    # Sets global seed for reproducibility.
    seed_everything(seed=args.seed, workers=True)

    datamodule = datamodule_or_datamodule_class
    if isclass(datamodule):
        datamodule = datamodule(args)

    # Sets persistent workers when using ddp_spawn. Must be set here because a
    # DataModule may be passed as an arg to main prior to ddp initialization.
    if args.devices > 1 and \
       "ddp_spawn" in args.strategy and \
       not datamodule.persistent_workers:
        datamodule.persistent_workers = True

    args.num_classes = datamodule.num_classes
    args.num_groups = datamodule.num_groups

    model = model_or_model_class
    if isclass(model):
        model = model(args)
        
    model = load_weights(args, model)

    if model_hooks:
        for hook in model_hooks:
            hook(model)

    trainer = load_trainer(args, addtl_callbacks=callbacks)
    if not args.eval_only:
        trainer.fit(model, datamodule=datamodule, ckpt_path=args.ckpt_path)

    # Reloads trainer for validation and testing. Must use a single device.
    # Note that this procedure resets the "epoch" and "step" trainer metrics.
    orig_devices = args.devices
    orig_strategy = args.strategy
    if args.devices > 1 and trainer.is_global_zero:
        if "ddp_spawn" not in args.strategy:
            torch.distributed.destroy_process_group()
        args.devices = 1
        trainer = load_trainer(args, addtl_callbacks=callbacks)

    val_metrics = trainer.validate(model, datamodule=datamodule, verbose=verbose)
    test_metrics = trainer.test(model, datamodule=datamodule, verbose=verbose)
    args.devices = orig_devices
    args.strategy = orig_strategy

    # Closes wanbd instance (for experiments which run main() many times).
    # Also cleans up wandb cache by deleting files down to a 10GB limit.
    if args.wandb:
        wandb.finish()
        cache = wandb.sdk.artifacts.artifacts_cache.get_artifacts_cache()
        cache.cleanup(int(10e10))
    
    return model, val_metrics, test_metrics


if __name__ == "__main__":
    args = parse_args()
    
    models, datamodules = valid_models_and_datamodules()

    main(args, models[args.model], datamodules[args.datamodule])


milkshake/main.py
tmlabonte/last-layer-retraining
"""BERT model implementation."""

# Imports Python builtins.
import types

# Imports Python packages.
from transformers import BertForSequenceClassification, get_scheduler

# Imports PyTorch packages.
from torch.optim import SGD

# Imports milkshake packages.
from milkshake.models.model import Model


class BERT(Model):
    """BERT model implementation."""

    def __init__(self, args):
        """Initializes a BERT model.

        Args:
            args: The configuration dictionary.
        """

        super().__init__(args)

        self.model = BertForSequenceClassification.from_pretrained(
            "bert-base-uncased",
            num_labels=args.num_classes,
        )

        self.model.fc = self.model.classifier
        delattr(self.model, "classifier")

        def classifier(self, x):
            return self.fc(x)
        
        self.model.classifier = types.MethodType(classifier, self.model)
        self.model.base_forward = self.model.forward

        def forward(self, x):
            return self.base_forward(
                input_ids=x[:, :, 0],
                attention_mask=x[:, :, 1],
                token_type_ids=x[:, :, 2]).logits

        self.model.forward = types.MethodType(forward, self.model)

        # Freezes all parameters except those in the last layer.
        if args.train_fc_only:
            for p in self.model.parameters():
                p.requires_grad = False
            for p in self.model.fc.parameters():
                p.requires_grad = True

    def load_msg(self):
        return f"Loading BERT Base Uncased pretrained on Book Corpus and English Wikipedia."

    def configure_optimizers(self):
        """Returns the optimizer and learning rate scheduler."""

        no_decay = ["bias", "LayerNorm.weight"]
        decay_params = []
        nodecay_params = []
        for n, p in self.model.named_parameters():
            if not any(nd in n for nd in no_decay):
                decay_params.append(p)
            else:
                nodecay_params.append(p)

        optimizer_grouped_parameters = [
            {
                "params": decay_params,
                "weight_decay": self.hparams.weight_decay,
            },
            {
                "params": nodecay_params,
                "weight_decay": 0.0,
            },
        ]

        optimizer = self.optimizer(
            optimizer_grouped_parameters,
            lr=self.hparams.lr,
        )
        if isinstance(optimizer, SGD):
            optimizer.momentum = self.hparams.momentum

        # Only the linear scheduler and no scheduler are currently implemented.
        if self.hparams.lr_scheduler == "linear":
            scheduler = get_scheduler(
                "linear",
                optimizer=optimizer,
                num_warmup_steps=self.hparams.lr_warmup_epochs,
                num_training_steps=self.hparams.max_epochs,
            )

            return [optimizer], [scheduler]
        elif self.hparams.lr_scheduler == "step" and self.hparams.lr_steps == []:
            return optimizer
        else:
            raise NotImplementedError()


milkshake/models/bert.py
tmlabonte/last-layer-retraining
"""Parent class and training logic for a classification model."""

# Imports Python builtins.
from abc import abstractmethod
import numpy as np

# Imports PyTorch packages.
from pl_bolts.optimizers.lr_scheduler import LinearWarmupCosineAnnealingLR
import pytorch_lightning as pl
import torch
from torch import nn
import torch.nn.functional as F
from torch.nn.parameter import is_lazy
from torch.optim import Adam, AdamW, SGD
from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, MultiStepLR
import torchvision.models as models

# Imports milkshake packages.
from milkshake.utils import compute_accuracy, to_np

# TODO: Organize logging code in a separate file.
# TODO: Implement after_n_steps logging similarly to after_epoch.
# TODO: Clean up collate_metrics.


class Model(pl.LightningModule):
    """Parent class and training logic for a classification model.

    Attributes:
        self.hparams: The configuration dictionary.
        self.model: A torch.nn.Module.
        self.optimizer: A torch.optim optimizer.
    """

    def __init__(self, args):
        """Initializes a Model.

        Args:
            args: The configuration dictionary.
        """

        super().__init__()

        # Saves args into self.hparams.
        self.save_hyperparameters(args)
        print(self.load_msg())
        
        self.model = None
        
        optimizers = {"adam": Adam, "adamw": AdamW, "sgd": SGD}
        self.optimizer = optimizers[args.optimizer]

    @abstractmethod
    def load_msg(self):
        """Returns a descriptive message about the Model configuration."""

    def has_uninitialized_params(self):
        """Returns whether the model has uninitialized parameters."""

        for param in self.model.parameters():
            if is_lazy(param):
                return True
        return False

    def setup(self, stage):
        """Initializes model parameters if necessary.

        Args:
            stage: "train", "val", or "test".
        """

        if self.has_uninitialized_params():
            match stage:
                case "fit":
                    dataloader = self.trainer.datamodule.train_dataloader()
                case "validate":
                    dataloader = self.trainer.datamodule.val_dataloader()
                case "test":
                    dataloader = self.trainer.datamodule.test_dataloader()
                case "predict":
                    dataloader = self.trainer.datamodule.predict_dataloader()
            dummy_batch = next(iter(dataloader))
            self.forward(dummy_batch[0])

    def forward(self, inputs):
        """Predicts using the model.

        Args:
            inputs: A torch.Tensor of model inputs.
        
        Returns:
            The model prediction as a torch.Tensor.
        """
        return torch.squeeze(self.model(inputs), dim=-1)

    def configure_optimizers(self):
        """Returns the optimizer and learning rate scheduler."""

        optimizer = self.optimizer(
            self.model.parameters(),
            lr=self.hparams.lr,
            weight_decay=self.hparams.weight_decay,
        )

        if isinstance(optimizer, SGD):
            optimizer.momentum = self.hparams.momentum

        if self.hparams.lr_scheduler == "cosine":
            scheduler = CosineAnnealingLR(
                optimizer,
                T_max=self.hparams.max_epochs,
            )
        elif self.hparams.lr_scheduler == "cosine_warmup":
            scheduler = LinearWarmupCosineAnnealingLR(
                optimizer,
                self.hparams.lr_warmup_epochs,
                self.hparams.max_epochs,
            )
        elif self.hparams.lr_scheduler == "linear":
            scheduler = LinearLR(
                optimizer,
                start_factor=1,
                end_factor=self.hparams.lr_drop,
                total_iters=self.hparams.max_epochs,
            )
        elif self.hparams.lr_scheduler == "step":
            scheduler = MultiStepLR(
                optimizer,
                self.hparams.lr_steps,
                gamma=self.hparams.lr_drop,
            )

        return [optimizer], [scheduler]
    
    def log_helper(self, names, values, add_dataloader_idx=False):
        """Compresses calls to self.log.

        Args:
            names: A list of metric names to log.
            values: A list of metric values to log.
            add_dataloader_idx: Whether to include the dataloader index in the name.
        """

        for idx, (name, value) in enumerate(zip(names, values)):
            self.log(
                name,
                value,
                on_step=(name in ("loss", "train_loss")),
                on_epoch=(name not in ("loss", "train_loss")),
                prog_bar=(name in ("train_acc", "val_loss", "val_acc")),
                sync_dist=True,
                add_dataloader_idx=add_dataloader_idx,
            )
    
    def log_helper2(self, names, values, dataloader_idx):
        """Calls log_helper as necessary for each DataLoader.
        
        Args:
           names: A list of metric names to log.
           values: A list of metric values to log.
           dataloader_idx: The index of the current DataLoader.
        """

        if dataloader_idx == 0:
            self.log_helper(names, values)

        try:
            # Errors if there is only 1 dataloader -- this means we
            # already logged it, so just pass.
            self.log_helper(names, values, add_dataloader_idx=True)
        except:
            pass

    def log_metrics(self, result, stage, dataloader_idx):
        """Logs metrics using the step results.
        
        Args:
            result: The output of self.step.
            stage: "train", "val", or "test".
            dataloader_idx: The index of the current dataloader.
        """

        names = []
        values = []
        for name in self.hparams.metrics:
            if name in result:
                # Class and group metrics must be manually collated.
                if "by_class" not in name and "by_group" not in name:
                    names.append(f"{stage}_{name}")
                    values.append(result[name])

        names.extend(["epoch", "step"])
        values.extend([float(self.current_epoch), float(self.trainer.global_step)])
        
        self.log_helper2(names, values, dataloader_idx)
    
    def collate_metrics(self, step_results, stage):
        """Collates and logs metrics by class and group.
        
        This is necessary because the logger does not utilize the info of how many samples
        from each class/group are in each batch, so logging class/group metrics on_epoch as
        usual will weight each batch the same instead of adjusting for totals.
        
        Args:
            step_results: List of dictionary results of self.validation_step or self.test_step.
            stage: "val" or "test".
        """
        
        def collate_and_sum(name):
            stacked = torch.stack([result[name] for result in step_results])
            return torch.sum(stacked, 0)

        dataloader_idx = step_results[0]["dataloader_idx"]

        if "acc_by_class" in self.hparams.metrics or \
          "acc5_by_class" in self.hparams.metrics or \
          "acc_by_group" in self.hparams.metrics or \
          "acc5_by_group" in self.hparams.metrics:
            names = []
            values = []
            total_by_class = collate_and_sum("total_by_class")
            total_by_group = collate_and_sum("total_by_group")

            if "acc_by_class" in self.hparams.metrics:
                acc_by_class = collate_and_sum("correct_by_class") / total_by_class
                names.extend([f"{stage}_acc_class{j}"
                              for j in range(len(acc_by_class))])
                values.extend(list(acc_by_class))
            if "acc5_by_class" in self.hparams.metrics:
                acc5_by_class = collate_and_sum("correct5_by_class") / total_by_class
                names.extend([f"{stage}_acc5_class{j}"
                              for j in range(len(acc5_by_class))])
                values.extend(list(acc5_by_class))
            if "acc_by_group" in self.hparams.metrics:
                acc_by_group = collate_and_sum("correct_by_group") / total_by_group
                names.extend([f"{stage}_acc_group{j}"
                              for j in range(len(acc_by_group))])
                values.extend(list(acc_by_group))
            if "acc5_by_group" in self.hparams.metrics:
                acc5_by_group = collate_and_sum("correct5_by_group") / total_by_group
                names.extend([f"{stage}_acc5_group{j}"
                              for j in range(len(acc5_by_group))])
                values.extend(list(acc5_by_group))

            self.log_helper2(names, values, dataloader_idx)
        
    def add_metrics_to_result(self, result, accs, dataloader_idx):
        """Adds dataloader_idx and metrics from compute_accuracy to result dict.
        
        Args:
            result: A dictionary containing the loss, prediction probabilities, and targets.
            accs: The output of compute_accuracy.
            dataloader_idx: The index of the current DataLoader.
        """
        
        result["dataloader_idx"] = dataloader_idx
        result["acc"] = accs["acc"]
        result["acc5"] = accs["acc5"]

        for k in ["class", "group"]:
            result[f"acc_by_{k}"] = accs[f"acc_by_{k}"]
            result[f"acc5_by_{k}"] = accs[f"acc5_by_{k}"]
            result[f"correct_by_{k}"] = accs[f"correct_by_{k}"]
            result[f"correct5_by_{k}"] = accs[f"correct5_by_{k}"]
            result[f"total_by_{k}"] = accs[f"total_by_{k}"]
    
    def step(self, batch, idx):
        """Performs a single step of prediction and loss calculation.

        Args:
            batch: A tuple containing the inputs and targets as torch.Tensor.
            idx: The index of the given batch.

        Returns:
            A dictionary containing the loss, prediction probabilities, and targets.
        
        Raises:
            ValueError: Class weights are specified with MSE loss, or MSE loss
            is specified for a multiclass classification task.
        """

        inputs, orig_targets = batch

        # Removes extra targets (e.g., group index used for metrics).
        if orig_targets[0].ndim > 0:
            targets = orig_targets[:, 0]
        else:
            targets = orig_targets

        logits = self(inputs)

        # Ensures logits is a torch.Tensor.
        if isinstance(logits, (tuple, list)):
            logits = torch.squeeze(logits[0], dim=-1)

        # Initializes class weights if desired.
        weights = torch.ones(self.hparams.num_classes, device=logits.device)
        if self.hparams.class_weights:
            if self.hparams.loss == "mse":
                raise ValueError("Cannot use class weights with MSE.")
            weights = torch.Tensor(self.hparams.class_weights, device=logits.device)

        # Computes loss and prediction probabilities.
        if self.hparams.loss == "cross_entropy":
            if self.hparams.num_classes == 1:
                loss = F.binary_cross_entropy_with_logits(logits, targets, weight=weights)
                probs = torch.sigmoid(logits)
            else:
                loss = F.cross_entropy(logits, targets, weight=weights)
                probs = F.softmax(logits, dim=1)
        elif self.hparams.loss == "mse":
            if self.hparams.num_classes == 1:
                loss = F.mse_loss(logits, targets.float())
            elif self.hparams.num_classes == 2:
                loss = F.mse_loss(logits[:, 0], targets.float())
            else:
                raise ValueError("MSE is only an option for binary classification.")

        return {"loss": loss, "probs": probs, "targets": orig_targets}
    
    def step_and_log_metrics(self, batch, idx, dataloader_idx, stage):
        """Performs a step, then computes and logs metrics.
        
        Args:
            batch: A tuple containing the inputs and targets as torch.Tensor.
            idx: The index of the given batch.
            dataloader_idx: The index of the current dataloader.
            stage: "train", "val", or "test".

        Returns:
            A dictionary containing the loss, prediction probabilities, targets, and metrics.
        """
        
        result = self.step(batch, idx)
        
        accs = compute_accuracy(
            result["probs"],
            result["targets"],
            self.hparams.num_classes,
            self.hparams.num_groups,
        )
        
        self.add_metrics_to_result(result, accs, dataloader_idx)

        self.log_metrics(result, stage, dataloader_idx)

        return result

    def training_step(self, batch, idx, dataloader_idx=0):
        """Performs a single training step.

        Args:
            batch: A tuple containing the inputs and targets as torch.Tensor.
            idx: The index of the given batch.
            dataloader_idx: The index of the current dataloader.

        Returns:
            A dictionary containing the loss, prediction probabilities, targets, and metrics.
        """
        
        return self.step_and_log_metrics(batch, idx, dataloader_idx, "train")

    def training_epoch_end(self, training_step_outputs):
        """Collates metrics upon completion of the training epoch.
        
        Args:
            training_step_outputs: List of dictionary outputs of self.training_step.
        """
        
        self.collate_metrics(training_step_outputs, "train")

    
    def validation_step(self, batch, idx, dataloader_idx=0):
        """Performs a single validation step.

        Args:
            batch: A tuple containing the inputs and targets as torch.Tensor.
            idx: The index of the given batch.
            dataloader_idx: The index of the current dataloader.

        Returns:
            A dictionary containing the loss, prediction probabilities, targets, and metrics.
        """

        return self.step_and_log_metrics(batch, idx, dataloader_idx, "val")
    
    def validation_epoch_end(self, validation_step_outputs):
        """Collates metrics upon completion of the validation epoch.
        
        Args:
            validation_step_outputs: List of dictionary outputs of self.validation_step.
        """
        
        self.collate_metrics(validation_step_outputs, "val")

    def test_step(self, batch, idx, dataloader_idx=0):
        """Performs a single test step.

        Args:
            batch: A tuple containing the inputs and targets as torch.Tensor.
            idx: The index of the given batch.
            dataloader_idx: The index of the current dataloader.

        Returns:
            A dictionary containing the loss, prediction probabilities, targets, and metrics.
        """

        return self.step_and_log_metrics(batch, idx, dataloader_idx, "test")
    
    def test_epoch_end(self, test_step_outputs):
        """Collates metrics upon completion of the test epoch.
        
        Args:
            test_step_outputs: List of dictionary outputs of self.test_step.
        """
        
        self.collate_metrics(test_step_outputs, "test")

    def predict_step(self, batch, batch_idx, dataloader_idx=0):
        """Performs a single prediction step.

        Args:
            batch: A tuple containing the inputs and targets as torch.Tensor.
            idx: The index of the given batch.
            dataloader_idx: The index of the current dataloader.
        """

        logits = self(batch)

        if self.hparams.num_classes == 1:
            probs = torch.sigmoid(logits)
            preds = (probs >= 0.5).int()
        else:
            probs = F.softmax(logits, dim=1)
            preds = torch.argmax(probs, dim=1)

        return preds


milkshake/models/model.py
tmlabonte/last-layer-retraining
"""ResNet model implementation."""

# Imports PyTorch packages.
import torch
from torch import nn
import torchvision.models as models

# Imports milkshake packages.
from milkshake.models.model import Model


class ResNet(Model):
    """ResNet model implementation."""

    def __init__(self, args):
        """Initializes a ResNet model.

        Args:
            args: The configuration dictionary.
        """

        super().__init__(args)

        resnets = {
            18: models.resnet18,
            34: models.resnet34,
            50: models.resnet50,
            101: models.resnet101,
            152: models.resnet152,
        }

        # TODO: Include different ResNet pretraining options.
        if args.resnet_pretrained:
            self.model = resnets[args.resnet_version](weights="IMAGENET1K_V1")
        else:
            self.model = resnets[args.resnet_version](weights=None)

        # Reduces the kernel size and stride for smaller inputs, e.g., CIFAR-10
        # images (32 x 32) instead of ImageNet images (224 x 224).
        if args.resnet_small_input:
            self.model.conv1 = nn.Conv2d(
                args.input_channels,
                64,
                kernel_size=3,
                stride=1,
                padding=1,
                bias=False,
            )
            self.model.maxpool = nn.Identity()
        elif args.input_channels != 3:
            self.model.conv1 = nn.Conv2d(
                args.input_channels,
                64,
                kernel_size=7,
                stride=2,
                padding=3,
                bias=False,
            )

        self.model.fc = nn.Sequential(
            nn.Dropout(p=args.dropout_prob),
            nn.Linear(self.model.fc.in_features, args.num_classes),
        )

        # Freezes all parameters except those in the last layer.
        if args.train_fc_only:
            for p in self.model.parameters():
                p.requires_grad = False
            for p in self.model.fc.parameters():
                p.requires_grad = True

    def load_msg(self):
        if self.hparams.resnet_pretrained:
            return f"Loading ImageNet1K-pretrained ResNet{self.hparams.resnet_version}."
        else:
            return f"Loading ResNet{self.hparams.resnet_version} with no pretraining."

    def step(self, batch, idx):
        """Performs a single step of prediction and loss calculation.

        Args:
            batch: A tuple containing the inputs and targets as torch.tensor.
            idx: The index of the given batch.

        Returns:
            A dictionary containing the loss, prediction probabilities, and targets.
            The probs and targets are moved to CPU to free up GPU memory.
        """

        result = super().step(batch, idx)

        # Optionally adds regularization penalizing the l1 norm of model parameters.
        if self.hparams.resnet_l1_regularization:
            if self.hparams.train_fc_only:
                params = torch.cat([param.view(-1) for param in self.model.fc.parameters()])
            else:
                params = torch.cat([param.view(-1) for param in self.model.parameters()])

            param_l1_norm = torch.linalg.vector_norm(params, ord=1)
            result["loss"] += self.hparams.resnet_l1_regularization * param_l1_norm

        return result


milkshake/models/resnet.py
tmlabonte/last-layer-retraining
"""Utility functions for milkshake."""

# Imports Python packages.
import os
import os.path as osp
from glob import glob
import math
import numpy as np
import warnings

# Imports PyTorch packages.
import torch
from torch._utils import _accumulate 

# Imports milkshake packages.
from milkshake.datamodules.dataset import Subset


def get_weights(args, version, best=None, idx=None):
    """Returns weights path given model version and checkpoint index.
    
    Args:
        args: The configuration dictionary.
        version: The model's PyTorch Lightning version.
        best: Whether to return the best model checkpoint.
        idx: The model's checkpoint index (-1 selects the latest checkpoint).
    
    Returns:
        The filepath of the desired model weights.
    """

    ckpt_path = ""
    if args.wandb:
        ckpt_path = args.wandb_dir
    
    # Selects the right naming convention for PL versions based on
    # whether the version input is an int or a string.
    if type(version) == int:
        ckpt_path = osp.join(ckpt_path, f"lightning_logs/version_{version}/checkpoints/*")
    else:
        ckpt_path = osp.join(ckpt_path, f"lightning_logs/{version}/checkpoints/*")

    list_of_weights = glob(osp.join(os.getcwd(), ckpt_path))
    
    if best:
        return [w for w in list_of_weights if "best" in w][0]
    else:
        list_of_weights = sorted([w for w in list_of_weights if "best" not in w])
        return list_of_weights[idx]

def compute_accuracy(probs, targets, num_classes, num_groups):
    """Computes top-1 and top-5 accuracies.
    
    Computes top-1 and top-5 accuracies by total and by class, and also
    includes the number of correct predictions and number of samples.
    The latter is useful for, e.g., collating metrics over an epoch.
    If groups are provided (as a second column in targets), then
    also returns group accuracies.

    Args:
        probs: A torch.Tensor of prediction probabilities.
        targets: A torch.Tensor of classification targets.
        num_classes: The total number of classes.
        num_groups: The total number of groups.

    Returns:
        A dictionary of metrics including top-1 and top-5 accuracies, number of
        correct predictions, and number of samples, each by total, class, and group.
    """

    # TODO: Clean up group metrics.

    # Splits apart targets and groups if necessary.
    groups = None
    if num_groups:
        groups = targets[:, 1]
        targets = targets[:, 0]

    if num_classes == 1:
        preds1 = (probs >= 0.5).int()
    else:
        preds1 = torch.argmax(probs, dim=1)

    correct1 = (preds1 == targets).float()
    
    acc1_by_class = []
    correct1_by_class = []
    total_by_class = []
    for j in range(num_classes):
        correct_nums = correct1[targets == j]
        acc1_by_class.append(correct_nums.mean())
        correct1_by_class.append(correct_nums.sum())
        total_by_class.append((targets == j).sum())

    acc1_by_group = [torch.tensor(1., device=targets.device)]
    correct1_by_group = [torch.tensor(1., device=targets.device)]
    total_by_group = [torch.tensor(1., device=targets.device)]
    if num_groups:
        acc1_by_group = []
        correct1_by_group = []
        total_by_group = []
        for j in range(num_groups):
            correct_nums = correct1[groups == j]
            acc1_by_group.append(correct_nums.mean())
            correct1_by_group.append(correct_nums.sum())
            total_by_group.append((groups == j).sum())

    correct5 = torch.tensor([1.] * len(targets), device=targets.device)
    acc5_by_class = [torch.tensor(1., device=targets.device)] * num_classes
    correct5_by_class = total_by_class
    if num_classes > 5:
        _, preds5 = torch.topk(probs, k=5, dim=1)
        correct5 = torch.tensor([t in preds5[j] for j, t in enumerate(targets)])
        correct5 = correct5.float()

        acc5_by_class = []
        correct5_by_class = []
        for j in range(num_classes):
            correct_nums = correct5[targets == j]
            acc5_by_class.append(correct_nums.mean())
            correct5_by_class.append(correct_nums.sum())

    acc5_by_group = [torch.tensor(1., device=targets.device)]
    correct5_by_group = [torch.tensor(1., device=targets.device)]
    if num_groups:
        acc5_by_group = [torch.tensor(1., device=targets.device)] * num_groups
        correct5_by_group = total_by_group
        if num_classes > 5:
            acc5_by_group = []
            correct5_by_group = []
            for j in range(num_groups):
                correct_nums = correct5[groups == j]
                acc5_by_group.append(correct_nums.mean())
                correct5_by_group.append(correct_nums.sum())

    accs = {
        "acc": correct1.mean(),
        "acc5": correct5.mean(),
        "acc_by_class": torch.stack(acc1_by_class),
        "acc5_by_class": torch.stack(acc5_by_class),
        "acc_by_group": torch.stack(acc1_by_group),
        "acc5_by_group": torch.stack(acc5_by_group),
        "correct": correct1.sum(),
        "correct5": correct5.sum(),
        "correct_by_class": torch.stack(correct1_by_class),
        "correct5_by_class": torch.stack(correct5_by_class),
        "correct_by_group": torch.stack(correct1_by_group),
        "correct5_by_group": torch.stack(correct5_by_group),
        "total": len(targets),
        "total_by_class": torch.stack(total_by_class),
        "total_by_group": torch.stack(total_by_group),
    }
    
    return accs

def _to_np(x):
    """Converts torch.Tensor input to numpy array."""

    return x.cpu().detach().numpy()

def to_np(x):
    """Converts input to numpy array.

    Args:
        x: A torch.Tensor, np.ndarray, or list.

    Returns:
        The input converted to a numpy array.

    Raises:
        ValueError: The input cannot be converted to a numpy array.
    """

    if not len(x):
        return np.array([])
    elif isinstance(x, torch.Tensor):
        return _to_np(x)
    elif isinstance(x, (np.ndarray, list)):
        if isinstance(x[0], torch.Tensor):
            return _to_np(torch.tensor(x))
        else:
            return np.asarray(x)
    else:
        raise ValueError("Input cannot be converted to numpy array.")

def random_split(dataset, lengths, generator):
    """Random split function from PyTorch adjusted for milkshake.Subset.
    
    Args:
        dataset: The milkshake.Dataset to be randomly split.
        lengths: The lengths or fractions of splits to be produced.
        generator: The generator used for the random permutation.
    
    Returns:
        A list of milkshake.Subsets with the desired splits.
    
    Raises:
        ValueError: The sum of input lengths does not equal the length of the input dataset.
    """

    # Handles the case where lengths is a list of fractions.
    if math.isclose(sum(lengths), 1) and sum(lengths) <= 1:
        subset_lengths = []
        for i, frac in enumerate(lengths):
            if frac < 0 or frac > 1:
                raise ValueError(f"Fraction at index {i} is not between 0 and 1")
            n_items_in_split = int(
                math.floor(len(dataset) * frac)
            )
            subset_lengths.append(n_items_in_split)
        remainder = len(dataset) - sum(subset_lengths)

        # Adds 1 to all the lengths in round-robin fashion until the remainder is 0.
        for i in range(remainder):
            idx_to_add_at = i % len(subset_lengths)
            subset_lengths[idx_to_add_at] += 1
        lengths = subset_lengths
        for i, length in enumerate(lengths):
            if length == 0:
                print(f"Length of split at index {i} is 0. "
                      f"This might result in an empty dataset.")

    if sum(lengths) != len(dataset):    # type: ignore[arg-type]
        raise ValueError("Sum of input lengths does not equal the length of the input dataset!")

    indices = torch.randperm(sum(lengths), generator=generator).tolist() 
    return [Subset(dataset, indices[offset - length : offset])
            for offset, length in zip(_accumulate(lengths), lengths)]

def ignore_warnings():
    """Adds nuisance warnings to ignore list.

    Should be called before import pytorch_lightning.
    """

    warnings.filterwarnings(
        "ignore",
        message=r"The feature ([^\s]+) is currently marked under review",
    )

    warnings.filterwarnings(
        "ignore",
        message=r"In the future ([^\s]+)",
    )

    warnings.filterwarnings(
        "ignore",
        message=r"Lazy modules ([^\s]+)",
    )

    warnings.filterwarnings(
        "ignore",
        message=r"There is a wandb run already in progress ([^\s]+)",
    )

    original_filterwarnings = warnings.filterwarnings
    def _filterwarnings(*xargs, **kwargs):
        return original_filterwarnings(*xargs, **{**kwargs, "append": True})
    warnings.filterwarnings = _filterwarnings


milkshake/utils.py
tmlabonte/last-layer-retraining
"""Setuptools procedure for building the milkshake package."""

from setuptools import setup

with open("requirements.txt") as f:
    lines = f.read()
    required = lines.splitlines()

setup(
    author="Tyler LaBonte",
    author_email="tlabonte@gatech.edu",
    description="Quick and extendable experimentation with classification models",
    install_requires=required,
    name="milkshake",
    packages=["milkshake"],
    version="1.0.0",
)

setup.py
tmlabonte/last-layer-retraining