First soup!

environment.yml

@@ -0,0 +1,81 @@
+name: model_soups
+channels:
+  - pytorch
+  - defaults
+dependencies:
+  - astroid=2.6.6=py36h06a4308_0
+  - blas=1.0=mkl
+  - ca-certificates=2022.4.26=h06a4308_0
+  - certifi=2021.5.30=py36h06a4308_0
+  - cudatoolkit=11.0.221=h6bb024c_0
+  - dataclasses=0.8=pyh4f3eec9_6
+  - freetype=2.11.0=h70c0345_0
+  - intel-openmp=2022.0.1=h06a4308_3633
+  - isort=5.9.3=pyhd3eb1b0_0
+  - jpeg=9b=h024ee3a_2
+  - lazy-object-proxy=1.6.0=py36h27cfd23_0
+  - lcms2=2.12=h3be6417_0
+  - ld_impl_linux-64=2.38=h1181459_1
+  - libffi=3.3=he6710b0_2
+  - libgcc-ng=11.2.0=h1234567_1
+  - libpng=1.6.37=hbc83047_0
+  - libstdcxx-ng=11.2.0=h1234567_1
+  - libtiff=4.2.0=h85742a9_0
+  - libuv=1.40.0=h7b6447c_0
+  - libwebp-base=1.2.2=h7f8727e_0
+  - lz4-c=1.9.3=h295c915_1
+  - mccabe=0.6.1=py36_1
+  - mkl=2020.2=256
+  - mkl-service=2.3.0=py36he8ac12f_0
+  - mkl_fft=1.3.0=py36h54f3939_0
+  - mkl_random=1.1.1=py36h0573a6f_0
+  - ncurses=6.3=h7f8727e_2
+  - ninja=1.10.2=h06a4308_5
+  - ninja-base=1.10.2=hd09550d_5
+  - numpy=1.19.2=py36h54aff64_0
+  - numpy-base=1.19.2=py36hfa32c7d_0
+  - olefile=0.46=pyhd3eb1b0_0
+  - openjpeg=2.4.0=h3ad879b_0
+  - openssl=1.1.1o=h7f8727e_0
+  - pillow=8.3.1=py36h2c7a002_0
+  - pip=21.2.2=py36h06a4308_0
+  - pylint=2.9.6=py36h06a4308_1
+  - python=3.6.13=h12debd9_1
+  - pytorch=1.7.1=py3.6_cuda11.0.221_cudnn8.0.5_0
+  - readline=8.1.2=h7f8727e_1
+  - setuptools=58.0.4=py36h06a4308_0
+  - six=1.16.0=pyhd3eb1b0_1
+  - sqlite=3.38.3=hc218d9a_0
+  - tk=8.6.11=h1ccaba5_1
+  - toml=0.10.2=pyhd3eb1b0_0
+  - torchvision=0.8.2=py36_cu110
+  - typed-ast=1.4.3=py36h7f8727e_1
+  - typing-extensions=4.1.1=hd3eb1b0_0
+  - typing_extensions=4.1.1=pyh06a4308_0
+  - wheel=0.37.1=pyhd3eb1b0_0
+  - wrapt=1.12.1=py36h7b6447c_1
+  - xz=5.2.5=h7f8727e_1
+  - zlib=1.2.12=h7f8727e_2
+  - zstd=1.4.9=haebb681_0
+  - pip:
+    - charset-normalizer==2.0.12
+    - clip==0.1.0
+    - cycler==0.11.0
+    - ftfy==6.0.3
+    - idna==3.3
+    - importlib-resources==5.4.0
+    - kiwisolver==1.3.1
+    - matplotlib==3.3.4
+    - pandas==1.1.5
+    - pyparsing==3.0.9
+    - python-dateutil==2.8.2
+    - pytz==2022.1
+    - regex==2022.4.24
+    - requests==2.27.1
+    - tqdm==4.64.0
+    - urllib3==1.26.9
+    - wcwidth==0.2.5
+    - wget==3.2
+    - zipp==3.6.0
+prefix: /home/mitchnw/anaconda3/envs/model_soups
+

finetune.py

@@ -0,0 +1,196 @@
+import argparse
+import os
+import torch
+import clip
+import os
+from tqdm import tqdm
+import time
+from utils import ModelWrapper, maybe_dictionarize_batch, cosine_lr
+from zeroshot import zeroshot_classifier
+import torch
+from torchvision import transforms, datasets
+
+def parse_arguments():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--data-location",
+        type=str,
+        default=os.path.expanduser('~/data'),
+        help="The root directory for the datasets.",
+    )
+    parser.add_argument(
+        "--model-location",
+        type=str,
+        default=os.path.expanduser('~/ssd/checkpoints/soups'),
+        help="Where to download the models.",
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=256,
+    )
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=8,
+    )
+    parser.add_argument(
+        "--epochs",
+        type=int,
+        default=8,
+    )
+    parser.add_argument(
+        "--warmup-length",
+        type=int,
+        default=500,
+    )
+    parser.add_argument(
+        "--lr",
+        type=float,
+        default=2e-5,
+    )
+    parser.add_argument(
+        "--wd",
+        type=float,
+        default=0.1,
+    )
+    parser.add_argument(
+        "--model",
+        default='ViT-B/32',
+        help='Model to use -- you can try another like ViT-L/14'
+    )
+    parser.add_argument(
+        "--name",
+        default='finetune_cp',
+        help='Filename for the checkpoints.'
+    )
+    parser.add_argument(
+        "--timm-aug", action="store_true", default=False,
+    )
+    parser.add_argument(
+    "--checkpoint_path",
+    default=None,
+    help='Checkpoint path to load the model'
+    )
+
+    return parser.parse_args()
+
+if __name__ == '__main__':
+    args = parse_arguments()
+    DEVICE = 'cuda'
+
+
+    template = [lambda x : f"a photo generated by {x}."]
+
+
+    base_model, preprocess = clip.load(args.model, 'cuda', jit=False)
+
+
+    train_transforms = transforms.Compose([transforms.RandomRotation(30),
+                                        transforms.RandomResizedCrop(224),
+                                        transforms.RandomHorizontalFlip(),
+                                        transforms.ToTensor()])
+
+    test_transforms = transforms.Compose([transforms.RandomRotation(30),
+                                        transforms.RandomResizedCrop(224),
+                                        transforms.ToTensor()])
+
+
+    train_data = datasets.ImageFolder(args.data_location + '/train', transform=train_transforms)
+    test_data = datasets.ImageFolder(args.data_location + '/test', transform=test_transforms)
+
+    train_dset = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, num_workers=args.workers, shuffle = True)
+    test_dset = torch.utils.data.DataLoader(test_data, batch_size=args.batch_size, num_workers=args.workers)
+
+    clf = zeroshot_classifier(base_model, ['humans', 'AI'], template, DEVICE)
+    NUM_CLASSES = 2
+    feature_dim = base_model.visual.output_dim
+
+    model = ModelWrapper(base_model, feature_dim, NUM_CLASSES, normalize=True, initial_weights=clf, checkpoint_path = args.checkpoint_path)
+    for p in model.parameters():
+        p.data = p.data.float()
+
+    model = model.cuda()
+    devices = [x for x in range(torch.cuda.device_count())]
+    model = torch.nn.DataParallel(model,  device_ids=devices)
+
+    model_parameters = [p for p in model.parameters() if p.requires_grad]
+    optimizer = torch.optim.AdamW(model_parameters, lr=args.lr, weight_decay=args.wd)
+
+    num_batches = len(train_dset)
+    scheduler = cosine_lr(optimizer, args.lr, args.warmup_length, args.epochs * num_batches)
+
+    loss_fn = torch.nn.CrossEntropyLoss()
+
+    model_path = os.path.join(args.model_location, f'{args.name}.pt')
+    print('Saving model to', model_path)
+    torch.save(model.module.state_dict(), model_path)
+
+    for epoch in range(args.epochs):
+        # Train
+        model.train()
+        end = time.time()
+        for i, batch in enumerate(train_dset):
+            step = i + epoch * num_batches
+            scheduler(step)
+            optimizer.zero_grad()
+            batch = maybe_dictionarize_batch(batch)
+            inputs, labels = batch['images'].to(DEVICE), batch['labels'].to(DEVICE)
+            data_time = time.time() - end
+
+            logits = model(inputs)
+            loss = loss_fn(logits, labels)
+
+            loss.backward()
+
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+
+            optimizer.step()
+
+            batch_time = time.time() - end
+            end = time.time()
+
+            if i % 20 == 0:
+                percent_complete = 100.0 * i / len(train_dset)
+                print(
+                    f"Train Epoch: {epoch} [{percent_complete:.0f}% {i}/{len(train_dset)}]\t"
+                    f"Loss: {loss.item():.6f}\tData (t) {data_time:.3f}\tBatch (t) {batch_time:.3f}", flush=True
+                )
+
+        ## Evaluate
+        test_loader = test_dset
+        model.eval()
+
+        last_accuracy = 0.0
+
+        with torch.no_grad():
+            print('*'*80)
+            print('Starting eval')
+            correct, count = 0.0, 0.0
+            pbar = tqdm(test_loader)
+            for batch in pbar:
+                batch = maybe_dictionarize_batch(batch)
+                inputs, labels = batch['images'].to(DEVICE), batch['labels'].to(DEVICE)
+
+                logits = model(inputs)
+
+                loss = loss_fn(logits, labels)
+
+                pred = logits.argmax(dim=1, keepdim=True)
+                correct += pred.eq(labels.view_as(pred)).sum().item()
+                count += len(logits)
+                pbar.set_description(
+                    f"Val loss: {loss.item():.4f}   Acc: {100*correct/count:.2f}")
+            top1 = correct / count
+        print(f'Val acc at epoch {epoch}: {100*top1:.2f}')
+
+        curr_acc = 100*top1
+        if curr_acc > last_accuracy: 
+            print('Current acc: {}, Last acc: {}'.format(curr_acc, last_accuracy))
+            last_accuracy = curr_acc
+            model_path = os.path.join(args.model_location, f'{args.name}.pt')
+            print('Saving model to', model_path)
+            torch.save(model.module.state_dict(), model_path)
+        else: 
+            print('Not saving the model')
+

helper.bash

@@ -0,0 +1 @@
+python finetune.py --data-location /l/users/u21010238/data/AiorNot --model-location /home/sara.pieri/Documents/model-soups/models --batch-size 56 --name finetune_cp_AiorNot_model_0 --checkpoint_path /home/sara.pieri/Documents/model-soups/models/model_0.pt

utils.py

@@ -0,0 +1,140 @@
+import torch
+import math
+import time
+import numpy as np
+
+class ModelWrapper(torch.nn.Module):
+    def __init__(self, model, feature_dim, num_classes, normalize=False, initial_weights=None, checkpoint_path = None):
+        super(ModelWrapper, self).__init__()
+        self.model = model
+        self.classification_head = torch.nn.Linear(feature_dim, num_classes)
+        self.normalize = normalize
+        
+        if initial_weights is None:
+            initial_weights = torch.zeros_like(self.classification_head.weight)
+            torch.nn.init.kaiming_uniform_(initial_weights, a=math.sqrt(5))
+
+        self.classification_head.weight = torch.nn.Parameter(initial_weights.clone())
+        self.classification_head.bias = torch.nn.Parameter(torch.zeros_like(self.classification_head.bias))
+
+        # Note: modified. Get rid of the language part.
+        if hasattr(self.model, 'transformer'):
+            delattr(self.model, 'transformer')
+        
+        if checkpoint_path:
+            print("Loading checkpoint", checkpoint_path)
+            checkpoint = torch.load(checkpoint_path)
+            checkpoint.pop('classification_head.weight')
+            checkpoint.pop('classification_head.bias')
+            model.load_state_dict(checkpoint, strict=False)
+
+    def forward(self, images, return_features=False):
+        features = self.model.encode_image(images)
+        if self.normalize:
+            features = features / features.norm(dim=-1, keepdim=True)
+        logits = self.classification_head(features)
+        if return_features:
+            return logits, features
+        return logits
+
+def get_model_from_sd(state_dict, base_model):
+    feature_dim = state_dict['classification_head.weight'].shape[1]
+    num_classes = state_dict['classification_head.weight'].shape[0]
+    model = ModelWrapper(base_model, feature_dim, num_classes, normalize=True)
+    for p in model.parameters():
+        p.data = p.data.float()
+    model.load_state_dict(state_dict)
+    model = model.cuda()
+    devices = [x for x in range(torch.cuda.device_count())]
+    return torch.nn.DataParallel(model,  device_ids=devices)
+
+def maybe_dictionarize_batch(batch):
+    if isinstance(batch, dict):
+        return batch
+    if len(batch) == 2:
+        return {'images': batch[0], 'labels': batch[1]}
+    elif len(batch) == 3:
+        return {'images': batch[0], 'labels': batch[1], 'metadata': batch[2]}
+    else:
+        raise ValueError(f'Unexpected number of elements: {len(batch)}')
+
+
+def test_model_on_dataset(model, dataset):
+
+    model.eval()
+    device = 'cuda'
+    with torch.no_grad():
+        top1, correct, n = 0., 0., 0.
+        end = time.time()
+        loader = dataset.test_loader
+        if type(dataset).__name__ == 'ImageNet2p':
+            loader = dataset.train_loader
+            # assert to make sure the imagenet held-out minival logic is consistent across machines.
+            # tested on a few machines but if this fails for you please submit an issue and we will resolve.
+            assert dataset.train_dataset.__getitem__(dataset.sampler.indices[1000])['image_paths'].endswith('n01675722_4108.JPEG')
+
+        for i, batch in enumerate(loader):
+            batch = maybe_dictionarize_batch(batch)
+            inputs, labels = batch['images'].cuda(), batch['labels'].cuda()
+            data_time = time.time() - end
+            y = labels
+            if 'image_paths' in batch:
+                image_paths = batch['image_paths']
+
+            logits = model(inputs)
+
+            projection_fn = getattr(dataset, 'project_logits', None)
+            if projection_fn is not None:
+                logits = projection_fn(logits, device)
+
+            if hasattr(dataset, 'project_labels'):
+                y = dataset.project_labels(y, device)
+            if isinstance(logits, list):
+                logits = logits[0]
+
+
+            pred = logits.argmax(dim=1, keepdim=True).to(device)
+            if hasattr(dataset, 'accuracy'):
+                acc1, num_total = dataset.accuracy(logits, y, image_paths, None)
+                correct += acc1
+                n += num_total
+            else:
+                correct += pred.eq(y.view_as(pred)).sum().item()
+                n += y.size(0)
+
+            batch_time = time.time() - end
+            end = time.time()
+            if i % 20 == 0:
+                percent_complete = 100.0 * i / len(loader)
+                print(
+                    f"[{percent_complete:.0f}% {i}/{len(loader)}]\t"
+                    f"Acc: {100 * (correct/n):.2f}\tData (t) {data_time:.3f}\tBatch (t) {batch_time:.3f}"
+                )
+
+        top1 = correct / n
+        return top1
+
+
+def assign_learning_rate(param_group, new_lr):
+    param_group["lr"] = new_lr
+
+def _warmup_lr(base_lr, warmup_length, step):
+    return base_lr * (step + 1) / warmup_length
+
+def cosine_lr(optimizer, base_lrs, warmup_length, steps):
+    if not isinstance(base_lrs, list):
+        base_lrs = [base_lrs for _ in optimizer.param_groups]
+    assert len(base_lrs) == len(optimizer.param_groups)
+    def _lr_adjuster(step):
+        for param_group, base_lr in zip(optimizer.param_groups, base_lrs):
+            if step < warmup_length:
+                lr = _warmup_lr(base_lr, warmup_length, step)
+            else:
+                e = step - warmup_length
+                es = steps - warmup_length
+                lr = 0.5 * (1 + np.cos(np.pi * e / es)) * base_lr
+            assign_learning_rate(param_group, lr)
+    return _lr_adjuster
+
+
+

zeroshot.py

@@ -0,0 +1,23 @@
+import argparse
+import os
+import torch
+import clip
+import os
+from tqdm import tqdm
+
+def zeroshot_classifier(model, classnames, templates, device):
+    print('Building zero-shot classifier.')
+    with torch.no_grad():
+        zeroshot_weights = []
+        for classname in tqdm(classnames):
+            texts = [template(classname) for template in templates] #format with class
+            texts = clip.tokenize(texts).to(device) #tokenize
+            class_embeddings = model.encode_text(texts)
+            class_embeddings /= class_embeddings.norm(dim=-1, keepdim=True)
+            class_embedding = class_embeddings.mean(dim=0)
+            class_embedding /= class_embedding.norm()
+            zeroshot_weights.append(class_embedding)
+        zeroshot_weights = torch.stack(zeroshot_weights, dim=1).to(device)
+    return 100*zeroshot_weights.t()
+
+