Self-Supervision with FastAI
A tutorial of rotation-based self-supervision using FastAI2 & PyTorch!
- Introduction
- Experiment Layout
- FastAI Vision Model Creation Methods
- PyTorch Rotation/Classification Self-Supervised Dataset
- Rotation Prediction Data
- FastAI Vision Learner [Rotation]
- Original Classification Data
- FastAI Vision Learner [Transfer-Classification]
- FastAI Vision Learner [From Sratch-Classification]
- Conclusion
Introduction
This notebook is an introduction to self-supervised learning. In short, self-supervised learning has 2 components:
- Pretrain on a pretext task, where the labels can come from the data itself!
- Transfer the features, and train on the actual classification labels!
"What if we can get labels for free for unlabelled data and train unsupervised dataset in a supervised manner? We can achieve this by framing a supervised learning task in a special form to predict only a subset of information using the rest. In this way, all the information needed, both inputs and labels, has been provided. This is known as self-supervised learning." - Lilian Weng
Using FastAI2, we'll use rotation as a pretext task for learning representations/features of our data.
Here are some great overviews of self-supervised learning that I've come across:
In this notebook, we will be using the MNIST dataset.
Also check out ImageWang from FastAI themselves! It's a dataset designed for self-supervision tasks!
-
Train a model on a rotation prediction task.
- We will use all the training data for rotation prediction.
- Input: A rotated image.
- Target/Label: Classify the amount of degrees rotated.
- Our model should learn useful features that can transfer well for a classification task.
- (The model should learn what digits look like in order to be able to successfully predict the amount of rotation).
-
Transfer our rotation pretraining features to solve the classification task with much fewer labels, < 1% of the original data.
- Input: A normal image.
- Target/Label: The images' original categorical label.
- Classification accuracy should be decent, even with only using < 1% of the original data.
-
Train a classifier from scratch on the same amount of data used in experiment 2.
- Input: A normal image.
- Target/Label: The images' original categorical label.
- Classification accuracy should be low (lack of transfer learning & too few labeled data!)
- Model may overfit.
Warning: This Jupyter notebook runs with fastai2! Make sure you have it installed, use the cell below to install it :)
!pip install fastai --upgrade
# Uncomment and run the below line to get a fresh install of fastai, if needed
# !pip install fastai --upgrade
Important: Pay attention! It’s important. We will be using a small ConvNet to test our self-supervised learning method. The architecture is defined below in
Note that simple_arch.
simple_arch takes in one argument, pretrained. This is to allow FastAI to pass pretrained=True or pretrained=False when creating the model body! Below are some use cases of when we would want pretrained=True or pretrained=False.
-
pretrained=False= For training a new model on our rotation prediction task. -
pretrained=True= For transferring the learnt features from our rotation task pretraining to solve a classification task. -
pretrained=False= For training a new model from scratch on the main classification task (no transfer learning).
from fastai.vision.all import *
def simple_arch(pretrained=False):
# Note that FastAI will automatically cut at pooling layer for the body!
model = nn.Sequential(
nn.Conv2d(1, 4, 3, 1),
nn.BatchNorm2d(4),
nn.ReLU(),
nn.Conv2d(4, 16, 3, 1),
nn.BatchNorm2d(16),
nn.ReLU(),
nn.Conv2d(16, 32, 3, 1),
nn.BatchNorm2d(32),
nn.AdaptiveAvgPool2d(1),
)
if (pretrained):
print("Loading pretrained model...")
pretrained_weights = torch.load(save_path/'rot_pretrained.pt')
print(model.load_state_dict(pretrained_weights))
return model
The follow below code snippets are examples of how FastAI creates CNNs. Every model will have a body and a head
body = create_body(arch=simple_arch, pretrained=False)
body
head = create_head(nf=32, n_out=8, lin_ftrs=[])
head
# Note that FastAI automatically determines nf for the head!
model = create_cnn_model(arch=simple_arch, pretrained=False, n_out=8, lin_ftrs=[])
model
import torchvision
tensorToImage = torchvision.transforms.ToPILImage()
imageToTensor = torchvision.transforms.ToTensor()
# Uncomment and run the below lines if torchvision has trouble downloading MNIST (in the next cell)
# !wget -P data/MNIST/raw/ http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz
# !wget -P data/MNIST/raw/ http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz
# !wget -P data/MNIST/raw/ http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz
# !wget -P data/MNIST/raw/ http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz
torchvision.datasets.MNIST('data/', download=True)
Below we define a dataset, here's the docstring:
A Dataset for Rotation-based Self-Supervision! Images are rotated clockwise.
-
file- MNIST processed .pt file. -
pct- percent of data to use -
classification- False=Use rotation labels. True=Use original classification labels.
class Custom_Dataset_MNIST():
'''
A Dataset for Rotation-based Self-Supervision! Images are rotated clockwise.
- file - MNIST processed .pt file.
- pct - percent of data to use
- classification - False=Use rotation labels. True=Use original classification labels.
'''
def __init__(self, file, pct, classification):
data = torch.load(file)
self.imgs = data[0]
self.labels = data[1]
self.pct = pct
self.classification = classification
slice_idx = int(len(self.imgs)*self.pct)
self.imgs = self.imgs[:slice_idx]
def __len__(self):
return len(self.imgs)
def __getitem__(self, idx):
img = self.imgs[idx].unsqueeze(0)
img = tensorToImage(img)
img = img.resize((32, 32), resample=1)
img = imageToTensor(img)
if (not self.classification):
# 4 classes for rotation
degrees = [0, 45, 90, 135, 180, 225, 270, 315]
rand_choice = random.randint(0, len(degrees)-1)
img = tensorToImage(img)
img = img.rotate(degrees[rand_choice])
img = imageToTensor(img)
return img, torch.tensor(rand_choice).long()
return img, self.labels[idx]
def show_batch(self, n=3):
fig, axs = plt.subplots(n, n)
fig.tight_layout()
for i in range(n):
for j in range(n):
rand_idx = random.randint(0, len(self)-1)
img, label = self.__getitem__(rand_idx)
axs[i, j].imshow(tensorToImage(img), cmap='gray')
if self.classification:
axs[i, j].set_title('Label: {0} (Digit #{1})'.format(label.item(), label.item()))
else:
axs[i, j].set_title('Label: {0} ({1} Degrees)'.format(label.item(), label.item()*45))
axs[i, j].axis('off')
Important: 60k training data and 10k validation data!
train_ds = Custom_Dataset_MNIST('data/MNIST/processed/training.pt', pct=1.0, classification=False)
valid_ds = Custom_Dataset_MNIST('data/MNIST/processed/test.pt', pct=1.0, classification=False)
print('{0} Training Samples | {1} Validation Samples'.format(len(train_ds), len(valid_ds)))
Note: Notice that our labels don’t correspond to digits! They correspond to the amount of degrees rotated! Specifically from this predefined set:
[0, 45, 90, 135, 180, 225, 270, 315]
from fastai.data.core import DataLoaders
dls = DataLoaders.from_dsets(train_ds, valid_ds).cuda()
# Override the show_batch function of dls to the one used in our dataset!
dls.show_batch = train_ds.show_batch
# We have 8 classes! [0, 1, 2, 3, 4, 5, 6, 7] that correspond to the [0, 45, 90, 135, 180, 225, 270, 315] degrees of rotation.
dls.c = 8
dls.show_batch()
rotation_head = create_head(nf=32, n_out=8, lin_ftrs=[])
rotation_head
Note: We want to measure
top_2_accuracy along with regular (top_1) accuracy, because there are hard-cases where it’s understandable why our model got it wrong. For example: ’0’ rotated 90 or 270 degrees, or ’1’ rotated 0 or 180 degrees. (They can look the same!)
# - A zero rotated 90 or 270 degrees?
# - A one rotated 0 or 180 degrees?
# etc :P
top_2_accuracy = lambda inp, targ: top_k_accuracy(inp, targ, k=2)
top_2_accuracy
Here, we train a model on the rotation prediction task!
# Note to set a value for lin_ftrs, we use the defined config above.
learner = cnn_learner(dls,
simple_arch,
pretrained=False,
loss_func=CrossEntropyLossFlat(),
custom_head=rotation_head,
metrics=[accuracy, top_2_accuracy])
learner.model
learner.summary()
learner.lr_find()
learner.fit_one_cycle(5, lr_max=3e-2)
| epoch | train_loss | valid_loss | accuracy | | time |
</tr>
</thead>
0 |
0.932141 |
1.084204 |
0.502600 |
0.875300 |
00:12 |
1 |
0.861757 |
0.831546 |
0.631500 |
0.905800 |
00:12 |
2 |
0.757740 |
0.841575 |
0.612800 |
0.932800 |
00:12 |
3 |
0.678802 |
0.655642 |
0.706700 |
0.949000 |
00:11 |
4 |
0.629287 |
0.550476 |
0.762000 |
0.968400 |
00:12 |
Important: We were able to achieve 76.2% top-1 accuracy, and 96.8% top-2 accuracy after just 5 epochs! Now we want to grab our
model from our Learner, and save the body of it!
Note: Our
model has two components, the body and the head. model is a list of size 2, where model[0] is the body, and model[1] is the head!
trained_body = learner.model[0]
trained_body
Tip: To save a model in PyTorch, save it’s
state_dict function! You can use model.load_state_dict to re-load the weights.
save_path = Path('rotation_cps/')
if not save_path.exists():
save_path.mkdir()
# Save the rotation-pretraining weights of our model body
torch.save(trained_body.state_dict(), save_path/'rot_pretrained.pt')
Now that we have pretrained our model on the rotation prediction task, we want to switch over to the original labeled data for the classification task.
Important: We’re only using 180 samples for training!
# Use 100% classification labeled data for validation!
train_ds = Custom_Dataset_MNIST('data/MNIST/processed/training.pt', pct=0.003, classification=True)
valid_ds = Custom_Dataset_MNIST('data/MNIST/processed/test.pt', pct=1.0, classification=True)
print('{0} Training Samples | {1} Validation Samples'.format(len(train_ds), len(valid_ds)))
Note: Notice the labels now correspond to the digit class!
from fastai.data.core import DataLoaders
dls = DataLoaders.from_dsets(train_ds, valid_ds).cuda()
dls.show_batch = train_ds.show_batch
# We have 10 classes! One for each digit label!
dls.c = 10
dls.show_batch()
Here we will toggle classification_head = create_head(nf=32, n_out=10, lin_ftrs=[])
classification_head
Note: We have
n_out=10 because of the 10 different digit classes
# pretrained=True will load the saved rotation pretraining weights into our model's body!
# See simple_arch() function definition for more details!
learner = cnn_learner(dls,
simple_arch,
pretrained=True,
loss_func=CrossEntropyLossFlat(),
custom_head=classification_head,
metrics=[accuracy])
learner.model
Tip: Freezing a model’s body after transferring the weights over, allows the new head to get calibrated with the rest of the model!
learner.freeze()
Note: Looking at the model summary, we can see that the model is frozen up to the new head! Good!
learner.summary()
learner.lr_find()
learner.fit_one_cycle(10, lr_max=3e-2)
Tip: Unfreeze the model after calibrating the new head with the transferred body, and train a little more!
learner.unfreeze()
learner.summary()
learner.lr_find()
learner.fine_tune(5, base_lr=3e-2)
Important: We were able to get 60.6% accuracy using transfer learning from our pretraining on the rotation prediction task!
Here we train a model from scratch on the original 180 labeled data. # pretrained=False, Create the same model as before, but without using the rotation pretraining weights!
learner = cnn_learner(dls,
simple_arch,
pretrained=False,
loss_func=CrossEntropyLossFlat(),
custom_head=classification_head,
metrics=[accuracy])
learner.model
learner.summary()
learner.lr_find()
learner.fit_one_cycle(10, lr_max=3e-2)
Important: We were able to only get 10.3% accuracy with training from scratch
Important: Using self-supervision can help learn features that can transfer to a down-stream task, such as classification! In this example, we used rotation predication as our pretext task for feature representation learning. Pretraining our model on rotation prediction prior to training for classification allowed us to achieve 60.6% accuracy, on just 0.3% of the labeled data (180 samples)! Training from scratch with the same amount of data yields an accuracy of 10.3%. The motivation for using self-supervised learning is the ability to train models with decent accuracy without the need of much labeled data!
Note: Be sure to try other self-supervised learning methods (or perhaps find your own!) and compete on the ImageWang Leadboard! How will model size, data difficultly, and dataset size (number of samples) affect self-supervised learning?
|
|---|