run.py

import argparse
import torch
import torch.backends.cudnn as cudnn
from torchvision import models
from data_aug.contrastive_learning_dataset import ContrastiveLearningDataset
from models.resnet_simclr import ResNetSimCLR
from simclr import SimCLR

model_names = sorted(name for name in models.__dict__
                     if name.islower() and not name.startswith("__")
                     and callable(models.__dict__[name]))

parser = argparse.ArgumentParser(description='PyTorch SimCLR')
parser.add_argument('-data', metavar='DIR', default='./datasets',
                    help='path to dataset')
parser.add_argument('-dataset-name', default='stl10',
                    help='dataset name', choices=['stl10', 'cifar10', 'huggingface', 'local'])
parser.add_argument('--hf-dataset-id', default='tsbpp/fall2025_deeplearning',
                    help='Hugging Face dataset id to use when --dataset-name huggingface (e.g. tsbpp/fall2025_deeplearning)')
parser.add_argument('--hf-image-col', default='image',
                    help='Image column name in the Hugging Face dataset (default: image)')
parser.add_argument('--streaming', action='store_true',
                    help='Whether to stream the dataset from Hugging Face with no existing local copy')
parser.add_argument('-a', '--arch', metavar='ARCH', default='resnet18',
                    choices=model_names,
                    help='model architecture: ' +
                         ' | '.join(model_names) +
                         ' (default: resnet50)')
parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
                    help='number of data loading workers (default: 4)')
parser.add_argument('--epochs', default=200, type=int, metavar='N',
                    help='number of total epochs to run')
parser.add_argument('-b', '--batch-size', default=256, type=int,
                    metavar='N',
                    help='mini-batch size (default: 256), this is the total '
                         'batch size of all GPUs on the current node when '
                         'using Data Parallel or Distributed Data Parallel')
parser.add_argument('--lr', '--learning-rate', default=0.0003, type=float,
                    metavar='LR', help='initial learning rate', dest='lr')
parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
                    metavar='W', help='weight decay (default: 1e-4)',
                    dest='weight_decay')
parser.add_argument('--seed', default=None, type=int,
                    help='seed for initializing training. ')
parser.add_argument('--disable-cuda', action='store_true',
                    help='Disable CUDA')
parser.add_argument('--fp16-precision', action='store_true',
                    help='Whether or not to use 16-bit precision GPU training.')

parser.add_argument('--out_dim', default=128, type=int,
                    help='feature dimension (default: 128)')
parser.add_argument('--log-every-n-steps', default=100, type=int,
                    help='Log every n steps')
parser.add_argument('--temperature', default=0.07, type=float,
                    help='softmax temperature (default: 0.07)')
parser.add_argument('--n-views', default=2, type=int, metavar='N',
                    help='Number of views for contrastive learning training.')
parser.add_argument('--gpu-index', default=0, type=int, help='Gpu index.')
parser.add_argument('--subset_percent', default=100.0, type=float,
                    help='Percentage of the dataset to use for training (default: 100.0)')
parser.add_argument('--log-dir', default=None,
                    help='Directory to save logs and checkpoints')
parser.add_argument('--enable-wandb', action='store_true',
                    help='Enable logging to Weights & Biases')
parser.add_argument('--wandb-project', default='simclr',
                    help='Weights & Biases project name')
parser.add_argument('--wandb-entity', default=None,
                    help='Weights & Biases entity (username or team name)')


def main():
    args = parser.parse_args()
    assert args.n_views == 2, "Only two view training is supported. Please use --n-views 2."
    # check if gpu training is available
    if not args.disable_cuda and torch.cuda.is_available():
        args.device = torch.device('cuda')
        cudnn.deterministic = True
        cudnn.benchmark = True
    else:
        args.device = torch.device('cpu')
        args.gpu_index = -1

    dataset = ContrastiveLearningDataset(args.data)

    # if using Hugging Face dataset, also pass dataset id and image column
    if args.dataset_name == 'huggingface':
        train_dataset = dataset.get_dataset(args.dataset_name, args.n_views,
                                            hf_dataset_id=args.hf_dataset_id,
                                            hf_image_col=args.hf_image_col,
                                            streaming=args.streaming,
                                            subset_percent=args.subset_percent,
                                            seed=args.seed)
    elif args.dataset_name == 'local':
        train_dataset = dataset.get_dataset(args.dataset_name, args.n_views,
                                            subset_percent=args.subset_percent,
                                            seed=args.seed)
    else:
        train_dataset = dataset.get_dataset(args.dataset_name, args.n_views)

    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=args.batch_size, shuffle=(not args.streaming and args.dataset_name != 'local') or (args.dataset_name == 'local'), # shuffling is handled by .shuffle() for streaming
        num_workers=args.workers, pin_memory=True, drop_last=True)

    model = ResNetSimCLR(base_model=args.arch, out_dim=args.out_dim)

    if args.streaming:
        # For streaming datasets, we need to estimate the number of steps
        # The dataset has 500,000 images
        num_samples = 500_000
        if args.subset_percent < 100.0:
            num_samples = int(args.subset_percent / 100.0 * num_samples)
        total_steps = (num_samples // args.batch_size) * args.epochs
    else:
        total_steps = len(train_loader) * args.epochs


    optimizer = torch.optim.Adam(model.parameters(), args.lr, weight_decay=args.weight_decay)

    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=total_steps, eta_min=0,
                                                           last_epoch=-1)

    #  It’s a no-op if the 'gpu_index' argument is a negative integer or None.
    with torch.cuda.device(args.gpu_index):
        simclr = SimCLR(model=model, optimizer=optimizer, scheduler=scheduler, args=args)
        simclr.train(train_loader)


if __name__ == "__main__":
    main()