update to match the camera-ready version

Author: slackerjerry

__pycache__/cifar_mil_trainer.cpython-38.pyc

@@ -0,0 +1,87 @@
+import argparse
+import json
+from cifar_mil_trainer import *
+import os
+import torch
+import seaborn as sns
+import matplotlib.pyplot as plt
+
+def get_args():
+
+    parser = argparse.ArgumentParser(description='MNIST MIL benchmarks:')
+
+    parser.add_argument("--project_name", default="MNIST-MIL")
+    parser.add_argument('--wandb', default=False, type=bool)
+
+    # Model params
+    parser.add_argument('--mode', default="softmax", choices=["softmax", "entmax", "sparsemax", "gsh"])
+    parser.add_argument('--d_model', default=1024, type=int)
+    parser.add_argument('--input_size', default=3*32*32, type=int)
+    parser.add_argument('--model', default="pooling", type=str)
+    parser.add_argument('--num_pattern', default=20, type=int)
+    parser.add_argument('--n_heads', default=8, type=int)
+    parser.add_argument('--scale', default=0.01)
+    parser.add_argument('--update_steps', default=1, type=int)
+    parser.add_argument('--dropout', default=0.7, type=float)
+
+    # Training params
+    parser.add_argument('--lr', default=1e-3, type=float)
+    parser.add_argument('--epoch', default=100, type=int)
+    parser.add_argument('--seed', default=1111, type=int)
+
+    # Data params
+    parser.add_argument('--batch_size', default=64, type=int)
+    parser.add_argument('--train_size', default=10000, type=int)
+    parser.add_argument('--test_size', default=5000, type=int)
+    parser.add_argument('--pos_per_bag', default=1, type=int)
+    parser.add_argument('--bag_size', default=10, type=int)
+    parser.add_argument('--tgt_num', default=0, type=int)
+
+    args = parser.parse_args()
+
+    return vars(args)
+
+
+if __name__ == "__main__":
+
+    torch.set_num_threads(3)
+    config = get_args()
+    trails = 5
+    torch.manual_seed(config["seed"])
+
+    bag_size = config["bag_size"]
+    # bag_size = [5, 10, 20, 50, 100, 200, 300]
+    models = ["softmax", "sparsemax", "entmax", "gsh"]    
+    data_log = None
+
+    for m in models:
+        config["mode"] = m
+        for t in range(trails):
+            torch.random.manual_seed(torch.random.seed())
+            trainer = Trainer(config, t)
+            trail_log = trainer.train()
+            if data_log is None:
+                data_log = trail_log
+            else:
+                for k,v in data_log.items():
+                    data_log[k] = data_log[k] + trail_log[k]
+    
+    sns.lineplot(data=data_log, x="epoch", y="train loss", hue="model", alpha=0.4, errorbar=None, linewidth=2)
+    plt.tight_layout()
+    plt.savefig(f'./imgs/cifar/train_loss_{bag_size}.pdf')
+    plt.clf()
+
+    sns.lineplot(data=data_log, x="epoch", y="test loss", hue="model", alpha=0.4, errorbar=None, linewidth=2)
+    plt.tight_layout()
+    plt.savefig(f'./imgs/cifar/test_loss_{bag_size}.pdf')
+    plt.clf()
+
+    sns.lineplot(data=data_log, x="epoch", y="train acc", hue="model", alpha=0.4, errorbar=None, linewidth=2)
+    plt.tight_layout()
+    plt.savefig(f'./imgs/cifar/train_acc_{bag_size}.pdf')
+    plt.clf()
+
+    sns.lineplot(data=data_log, x="epoch", y="test acc", hue="model", alpha=0.4, errorbar=None, linewidth=2)
+    plt.tight_layout()
+    plt.savefig(f'./imgs/cifar/test_acc_{bag_size}.pdf')
+    plt.clf()

__pycache__/layers.cpython-38.pyc

@@ -0,0 +1,160 @@
+from datasets.cifar10_bags import CIFARBags
+from torch.utils.data import DataLoader
+from layers import *
+from models import *
+import wandb
+import pandas as pd
+
+class Trainer:
+
+    def __init__(self, config, trial) -> None:
+        self.config = config
+        self.trial = trial
+
+        if self.config["wandb"]:
+            run = wandb.init(
+                # Set the project where this run will be logged
+                project=self.config["project_name"] + " good",
+                # Track hyperparameters and run metadata
+                config=self.config)
+    
+    def _get_data(self):
+
+        trainset = CIFARBags(target_number=self.config["tgt_num"], 
+                                  bag_size=self.config["bag_size"], 
+                                  num_bag=self.config["train_size"],
+                                  pos_per_bag=self.config["pos_per_bag"],
+                                  seed=self.config["seed"],
+                                  train=True
+                                  )  
+
+        testset = CIFARBags(target_number=self.config["tgt_num"], 
+                                  bag_size=self.config["bag_size"], 
+                                  num_bag=self.config["test_size"],
+                                  pos_per_bag=self.config["pos_per_bag"],
+                                  seed=self.config["seed"],
+                                  train=False
+                                  )  
+
+        train_loader = DataLoader(trainset, batch_size=self.config["batch_size"], shuffle=True)
+        test_loader = DataLoader(testset, batch_size=self.config["batch_size"], shuffle=False)
+
+        return train_loader, test_loader
+
+    def _get_model(self):
+
+        model = MNISTModel(input_size=self.config["input_size"],
+                            d_model=self.config["d_model"],
+                            n_heads=self.config["n_heads"], 
+                            update_steps=self.config["update_steps"], 
+                            dropout=self.config["dropout"],
+                            mode=self.config["mode"],
+                            scale=self.config["scale"],
+                            num_pattern=self.config['num_pattern'])
+
+        return model.cuda()
+
+    def _get_opt(self):
+        return torch.optim.Adam(self.model.parameters(), lr=self.config["lr"], weight_decay=0.0)
+
+    def _get_cri(self):
+        return torch.nn.BCEWithLogitsLoss()
+
+    def test_epoch(self, loader):
+
+        total_loss = 0.0
+        total_cor, total_sample = 0, 0
+        total_step = 0
+
+        with torch.no_grad():
+            for x, y in loader:
+
+                total_sample += x.size(0)
+                total_step += 1
+                x, y = x.float().cuda(), y.float().cuda()
+                pred = self.model(x)
+                loss = self.cri(pred, y)
+
+                output = (pred>0.5).float()
+                total_cor += (output == y).float().sum()
+                total_loss += loss.item()
+        
+        return total_loss/total_step, total_cor/total_sample
+
+    def train_epoch(self, loader):
+
+        total_loss = 0.0
+        total_cor, total_sample = 0, 0
+        total_step = 0
+
+        for x, y in loader:
+
+            total_step += 1
+            total_sample += x.size(0)
+
+            self.opt.zero_grad()
+            x, y = x.float().cuda(), y.float().cuda()
+            pred = self.model(x)
+            loss = self.cri(pred, y)
+            loss.backward()
+            self.opt.step()
+
+            output = (pred>0.5).float()
+            total_cor += (output == y).float().sum()
+            total_loss += loss.item()
+        
+        return total_loss/total_step, total_cor/total_sample
+
+    def train(self):
+
+        train_loader, test_loader = self._get_data()
+        self.model = self._get_model()
+        self.opt = self._get_opt()
+        self.cri = self._get_cri()
+
+        best_test_acc = -1
+
+        data_log = {
+            'train loss':[],
+            'train acc':[],
+            'test loss':[],
+            'test acc':[],
+            'epoch':[],
+            'model':[]
+        }
+
+        self.sche = torch.optim.lr_scheduler.CosineAnnealingLR(self.opt, self.config["epoch"], eta_min=0, last_epoch=-1, verbose=False)
+
+        for epoch in range(1, self.config["epoch"]+1):
+
+            train_loss, train_acc = self.train_epoch(train_loader)
+            test_loss, test_acc = self.test_epoch(test_loader)
+            self.sche.step()
+
+            data_log['train loss'].append(train_loss)
+            data_log['test loss'].append(test_loss)
+            data_log['train acc'].append(train_acc.item())
+            data_log['test acc'].append(test_acc.item())
+            data_log['epoch'].append(epoch)
+            data_log['model'].append(self.config['mode'])
+
+            if test_acc >= best_test_acc:
+                best_test_acc = test_acc
+
+            if self.config["wandb"]:
+                wandb.log({
+                    "step": epoch,
+                    "train loss": train_loss,
+                    "train acc": train_acc.item()*100,
+                    "test loss": test_loss,
+                    "test acc": test_acc.item()*100
+                    }, step=epoch)
+                
+        if self.config["wandb"]:
+            wandb.log({"best test acc": best_test_acc})
+            wandb.log({"logs": data_log})
+        
+        if self.config["wandb"]:
+            wandb.finish()
+
+        return data_log

__pycache__/mnist_mil_trainer.cpython-38.pyc

@@ -0,0 +1,5 @@
+python3 cifar_mil_main.py --bag_size 20
+python3 cifar_mil_main.py --bag_size 50
+python3 cifar_mil_main.py --bag_size 5
+python3 cifar_mil_main.py --bag_size 10
+python3 cifar_mil_main.py --bag_size 2

__pycache__/models.cpython-38.pyc

@@ -0,0 +1,119 @@
+"""Pytorch dataset object that loads MNIST dataset as bags."""
+
+import numpy as np
+import torch
+import torch.utils.data as data_utils
+from torchvision import datasets, transforms
+import random
+
+
+class CIFARBags(data_utils.Dataset):
+    def __init__(
+            self,
+            target_number=9,
+            bag_size=10,
+            num_bag=500,
+            pos_per_bag=1,
+            seed=1,
+            train=True):
+        self.target_number = target_number
+        self.bag_size = bag_size
+        self.pos_per_bag = pos_per_bag
+        self.train = train
+        self.num_bag = num_bag
+
+        self.r = np.random.RandomState(seed)
+
+        self.num_in_train = 50000
+        self.num_in_test = 10000
+
+        if self.train:
+            self.train_bags_list, self.train_labels_list = self._create_bags()
+        else:
+            self.test_bags_list, self.test_labels_list = self._create_bags()
+
+    def _create_bags(self):
+
+        transform_train = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
+        ])
+
+        transform_test = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
+        ])
+
+        if self.train:
+            loader = data_utils.DataLoader(
+                datasets.CIFAR10(
+                    '../datasets',
+                    train=True,
+                    download=True,
+                    transform=transform_train),
+                batch_size=self.num_in_train,
+                shuffle=False)
+        else:
+            loader = data_utils.DataLoader(
+                datasets.CIFAR10(
+                    '../datasets',
+                    train=False,
+                    download=True,
+                    transform=transform_test),
+                batch_size=self.num_in_test,
+                shuffle=False)
+
+        for (batch_data, batch_labels) in loader:
+            all_imgs = batch_data
+            all_labels = batch_labels
+
+        bags_list = []
+        labels_list = []
+
+        pos_idx = [i for i, j in enumerate(
+            all_labels) if j == self.target_number]
+        neg_idx = [i for i, j in enumerate(
+            all_labels) if j != self.target_number]
+
+        pos_images = []
+        neg_images = []
+
+        for i, img in enumerate(all_imgs):
+            if all_labels[i] == self.target_number:
+                pos_images.append(img)
+            else:
+                neg_images.append(img)
+
+        self.all_pos_img = pos_images
+        self.all_neg_img = neg_images
+
+        for i in range(self.num_bag):
+
+            _pos_idx = random.sample(pos_idx,
+                                     self.pos_per_bag) + random.sample(neg_idx,
+                                                                       self.bag_size - self.pos_per_bag)
+            _neg_idx = random.sample(neg_idx, self.bag_size)
+            assert len(_pos_idx) == len(_neg_idx)
+
+            bags_list.append(all_imgs[_neg_idx])
+            labels_list.append(1)
+            bags_list.append(all_imgs[_pos_idx])
+            labels_list.append(0)
+
+        return bags_list, torch.tensor(labels_list)
+
+    def __len__(self):
+        if self.train:
+            return len(self.train_labels_list)
+        else:
+            return len(self.test_labels_list)
+
+    def __getitem__(self, index):
+        if self.train:
+            bag = self.train_bags_list[index]
+            label = self.train_labels_list[index]
+        else:
+            bag = self.test_bags_list[index]
+            label = self.test_labels_list[index]
+
+        return bag, label