Optimizing Average Precision Loss on Imbalanced CIFAR10 Dataset (SOAP)
Introduction
In this tutorial, you will learn how to quickly train a Resnet18
model by optimizing AUPRC with our novel APLoss
and SOAP
optimizer
[ref] on a binary image
classification task with CIFAR-10 dataset. After completion of this
tutorial, you should be able to use LibAUC to train your own models
on your own datasets.
Reference:
If you find this tutorial helpful in your work, please cite our library paper and the following papers:
@article{qi2021stochastic,
title={Stochastic Optimization of Areas Under Precision-Recall Curves with Provable Convergence},
author={Qi, Qi and Luo, Youzhi and Xu, Zhao and Ji, Shuiwang and Yang, Tianbao},
journal={Advances in Neural Information Processing Systems},
volume={34},
year={2021}}
Install LibAUC
Let’s start with installing our library here. In this tutorial, we will use the lastest version for LibAUC by using pip install -U
.
!pip install -U libauc
Importing LibAUC
Import required packages to use
from libauc.losses import APLoss
from libauc.optimizers import SOAP
from libauc.models import resnet18 as ResNet18
from libauc.datasets import CIFAR10
from libauc.utils import ImbalancedDataGenerator
from libauc.sampler import DualSampler
from libauc.metrics import auc_prc_score
import torchvision.transforms as transforms
from torch.utils.data import Dataset
import numpy as np
import torch
from PIL import Image
Reproducibility
The following function set_all_seeds
limits the number of sources
of randomness behaviors, such as model intialization, data shuffling,
etcs. However, completely reproducible results are not guaranteed
across PyTorch releases
[Ref].
def set_all_seeds(SEED):
# REPRODUCIBILITY
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
set_all_seeds(2023)
Loading datasets
In this step, we will use the
CIFAR10 as
benchmark dataset. Before importing data to dataloader
, we
construct imbalanced version for CIFAR10 by
ImbalanceDataGenerator
. Specifically, it first randomly splits
the training data by class ID (e.g., 10 classes) into two even
portions as the positive and negative classes, and then it randomly
removes some samples from the positive class to make it imbalanced.
We keep the testing set untouched. We refer imratio
to the ratio
of number of positive examples to number of all examples.
train_data, train_targets = CIFAR10(root='./data', train=True).as_array()
test_data, test_targets = CIFAR10(root='./data', train=False).as_array()
imratio = 0.02 ## we set the imratio as 0.02 here since AP metric is usually used to evaluate highly imbalanced data.
generator = ImbalancedDataGenerator(verbose=True, random_seed=2023)
(train_images, train_labels) = generator.transform(train_data, train_targets, imratio=imratio)
(test_images, test_labels) = generator.transform(test_data, test_targets, imratio=0.5)
We define the data input pipeline such as data
augmentations. In this tutorial, we use RandomCrop
,
RandomHorizontalFlip
.
class ImageDataset(Dataset):
def __init__(self, images, targets, image_size=32, crop_size=30, mode='train'):
self.images = images.astype(np.uint8)
self.targets = targets
self.mode = mode
self.transform_train = transforms.Compose([
transforms.ToTensor(),
transforms.RandomCrop((crop_size, crop_size), padding=None),
transforms.RandomHorizontalFlip(),
transforms.Resize((image_size, image_size)),
])
self.transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Resize((image_size, image_size)),
])
def __len__(self):
return len(self.images)
def __getitem__(self, idx):
image = self.images[idx]
target = self.targets[idx]
image = Image.fromarray(image.astype('uint8'))
if self.mode == 'train':
image = self.transform_train(image)
else:
image = self.transform_test(image)
return image, target, idx
We define dataset
, DualSampler
and dataloader
here. By
default, we use batch_size
64 and we oversample the minority
class with pos:neg=1:1
by setting sampling_rate=0.5
.
batch_size = 64
sampling_rate = 0.5
trainSet = ImageDataset(train_images, train_labels)
trainSet_eval = ImageDataset(train_images, train_labels,mode='test')
testSet = ImageDataset(test_images, test_labels, mode='test')
sampler = DualSampler(trainSet, batch_size, sampling_rate=sampling_rate)
trainloader = torch.utils.data.DataLoader(trainSet, batch_size=batch_size, sampler=sampler, num_workers=2)
trainloader_eval = torch.utils.data.DataLoader(trainSet_eval, batch_size=batch_size, shuffle=False, num_workers=2)
testloader = torch.utils.data.DataLoader(testSet, batch_size=batch_size, shuffle=False, num_workers=2)
Configuration
Hyper-Parameters
lr = 1e-3
margin = 0.6
gamma = 0.1
weight_decay = 1e-5
total_epoch = 60
decay_epoch = [30]
Model, Loss and Optimizer
model = ResNet18(pretrained=False, last_activation=None, num_classes=1)
model = model.cuda()
loss_fn = APLoss(data_len=len(trainSet), margin=margin, gamma=gamma)
optimizer = SOAP(model.parameters(), lr=lr, mode='adam', weight_decay=weight_decay)
Training
Now it’s time for training. And we evaluate Average Precision performance after every epoch.
print ('Start Training')
print ('-'*30)
train_log, test_log = [], []
for epoch in range(total_epoch):
if epoch in decay_epoch:
optimizer.update_lr(decay_factor=10)
model.train()
for idx, (data, targets, index) in enumerate(trainloader):
data, targets, index = data.cuda(), targets.cuda(), index.cuda()
y_pred = model(data)
y_prob = torch.sigmoid(y_pred)
loss = loss_fn(y_prob, targets, index)
optimizer.zero_grad()
loss.backward()
optimizer.step()
######***evaluation***####
# evaluation on training sets
model.eval()
train_pred_list, train_true_list = [], []
for i, data in enumerate(trainloader_eval):
train_data, train_targets, _ = data
train_data = train_data.cuda()
y_pred = model(train_data)
y_prob = torch.sigmoid(y_pred)
train_pred_list.append(y_prob.cpu().detach().numpy())
train_true_list.append(train_targets.cpu().detach().numpy())
train_true = np.concatenate(train_true_list)
train_pred = np.concatenate(train_pred_list)
train_ap = auc_prc_score(train_true, train_pred)
train_log.append(train_ap)
# evaluation on test sets
model.eval()
test_pred_list, test_true_list = [], []
for j, data in enumerate(testloader):
test_data, test_targets, _ = data
test_data = test_data.cuda()
y_pred = model(test_data)
y_prob = torch.sigmoid(y_pred)
test_pred_list.append(y_prob.cpu().detach().numpy())
test_true_list.append(test_targets.numpy())
test_true = np.concatenate(test_true_list)
test_pred = np.concatenate(test_pred_list)
val_ap = auc_prc_score(test_true, test_pred)
test_log.append(val_ap)
model.train()
print("epoch: %s, train_ap: %.4f, test_ap: %.4f, lr: %.4f"%(epoch, train_ap, val_ap, optimizer.lr ))
Start Training
------------------------------
epoch: 0, train_ap: 0.0691, test_ap: 0.6653, lr: 0.0010
epoch: 1, train_ap: 0.1129, test_ap: 0.6586, lr: 0.0010
epoch: 2, train_ap: 0.3026, test_ap: 0.6780, lr: 0.0010
epoch: 3, train_ap: 0.5028, test_ap: 0.7093, lr: 0.0010
epoch: 4, train_ap: 0.6216, test_ap: 0.7308, lr: 0.0010
epoch: 5, train_ap: 0.6408, test_ap: 0.7268, lr: 0.0010
epoch: 6, train_ap: 0.7405, test_ap: 0.7058, lr: 0.0010
epoch: 7, train_ap: 0.6512, test_ap: 0.7060, lr: 0.0010
epoch: 8, train_ap: 0.8028, test_ap: 0.7318, lr: 0.0010
epoch: 9, train_ap: 0.8302, test_ap: 0.7222, lr: 0.0010
epoch: 10, train_ap: 0.9328, test_ap: 0.7284, lr: 0.0010
epoch: 11, train_ap: 0.8124, test_ap: 0.7122, lr: 0.0010
epoch: 12, train_ap: 0.9302, test_ap: 0.7143, lr: 0.0010
epoch: 13, train_ap: 0.9184, test_ap: 0.7208, lr: 0.0010
epoch: 14, train_ap: 0.7231, test_ap: 0.7018, lr: 0.0010
epoch: 15, train_ap: 0.9508, test_ap: 0.7216, lr: 0.0010
epoch: 16, train_ap: 0.9490, test_ap: 0.7220, lr: 0.0010
epoch: 17, train_ap: 0.9803, test_ap: 0.7337, lr: 0.0010
epoch: 18, train_ap: 0.8972, test_ap: 0.7276, lr: 0.0010
epoch: 19, train_ap: 0.9372, test_ap: 0.7314, lr: 0.0010
epoch: 20, train_ap: 0.9827, test_ap: 0.7293, lr: 0.0010
epoch: 21, train_ap: 0.9640, test_ap: 0.7354, lr: 0.0010
epoch: 22, train_ap: 0.9200, test_ap: 0.7205, lr: 0.0010
epoch: 23, train_ap: 0.9293, test_ap: 0.7118, lr: 0.0010
epoch: 24, train_ap: 0.9944, test_ap: 0.7229, lr: 0.0010
epoch: 25, train_ap: 0.9644, test_ap: 0.7219, lr: 0.0010
epoch: 26, train_ap: 0.9943, test_ap: 0.7216, lr: 0.0010
epoch: 27, train_ap: 0.9577, test_ap: 0.7176, lr: 0.0010
epoch: 28, train_ap: 0.9932, test_ap: 0.7196, lr: 0.0010
epoch: 29, train_ap: 0.8064, test_ap: 0.6964, lr: 0.0010
Reducing learning rate to 0.00010 @ T=23430!
epoch: 30, train_ap: 0.9969, test_ap: 0.7329, lr: 0.0001
epoch: 31, train_ap: 0.9975, test_ap: 0.7337, lr: 0.0001
epoch: 32, train_ap: 0.9976, test_ap: 0.7125, lr: 0.0001
epoch: 33, train_ap: 0.9978, test_ap: 0.6998, lr: 0.0001
epoch: 34, train_ap: 0.9979, test_ap: 0.7038, lr: 0.0001
epoch: 35, train_ap: 0.9977, test_ap: 0.7043, lr: 0.0001
epoch: 36, train_ap: 0.9980, test_ap: 0.6990, lr: 0.0001
epoch: 37, train_ap: 0.9980, test_ap: 0.6987, lr: 0.0001
epoch: 38, train_ap: 0.9980, test_ap: 0.6866, lr: 0.0001
epoch: 39, train_ap: 0.9979, test_ap: 0.6968, lr: 0.0001
epoch: 40, train_ap: 0.9981, test_ap: 0.6578, lr: 0.0001
epoch: 41, train_ap: 0.9981, test_ap: 0.6903, lr: 0.0001
epoch: 42, train_ap: 0.9979, test_ap: 0.6868, lr: 0.0001
epoch: 43, train_ap: 0.9982, test_ap: 0.6668, lr: 0.0001
epoch: 44, train_ap: 0.9981, test_ap: 0.6625, lr: 0.0001
epoch: 45, train_ap: 0.9979, test_ap: 0.6624, lr: 0.0001
epoch: 46, train_ap: 0.9983, test_ap: 0.6550, lr: 0.0001
epoch: 47, train_ap: 0.9982, test_ap: 0.6956, lr: 0.0001
epoch: 48, train_ap: 0.9983, test_ap: 0.6835, lr: 0.0001
epoch: 49, train_ap: 0.9981, test_ap: 0.6769, lr: 0.0001
epoch: 50, train_ap: 0.9979, test_ap: 0.7200, lr: 0.0001
epoch: 51, train_ap: 0.9980, test_ap: 0.7048, lr: 0.0001
epoch: 52, train_ap: 0.9981, test_ap: 0.7029, lr: 0.0001
epoch: 53, train_ap: 0.9986, test_ap: 0.6717, lr: 0.0001
epoch: 54, train_ap: 0.9995, test_ap: 0.6088, lr: 0.0001
epoch: 55, train_ap: 0.9999, test_ap: 0.6549, lr: 0.0001
epoch: 56, train_ap: 1.0000, test_ap: 0.6620, lr: 0.0001
epoch: 57, train_ap: 1.0000, test_ap: 0.6461, lr: 0.0001
epoch: 58, train_ap: 1.0000, test_ap: 0.6634, lr: 0.0001
epoch: 59, train_ap: 1.0000, test_ap: 0.6334, lr: 0.0001
Visualization
Now, let’s see the learning curve for optimizing AUPRC on train and test sets.
import matplotlib.pyplot as plt
plt.rcParams["figure.figsize"] = (9,5)
x=np.arange(len(train_log))
plt.figure()
plt.plot(x, train_log, lineStyle='-', label='APLoss Training', linewidth=3)
plt.plot(x, test_log, lineStyle='-', label='APLoss Test', linewidth=3)
plt.title('CIFAR-10 (2% imbalanced)',fontsize=25)
plt.legend(fontsize=15)
plt.ylabel('AP', fontsize=25)
plt.xlabel('Epoch', fontsize=25)
plt.show()

Comparison
Furthermore, we compare our library with TensorFlow Constrained Optimization (TFCO) library which can also be used to maximize AUPRC. For more information about TensorFlow Constrained Optimization (TFCO), please refer to this link. In the comparative experiment, we use same dataset, network and data augmentation. Although TFCO tutorial adopts Adagrad as the default optimizer, we also conduct experiment with Adam optimizer for fair comparison.