- 12.2, Fine-tuning an image captioning system.
So in our last session,
we talked about the Vision Transformer
and how it is revolutionizing computer vision,
the same way the Baseline Transformer
and its derivative architectures GPT, T5, BERT,
are revolutionizing natural language processing.
In this session, we're going to be fine-tuning
what is likely our most complicated system yet,
an image captioning system.
And to do that, we're actually going
to be putting together two transformers.
We'll be using the Vision Transformer
as our image processing unit.
And then combining that through cross-attention,
as we've seen with T5, with a DistilGPT2.
in order to caption the image, i.e.,
explain in a few words what is going on in that image.
Our main goal for our code segment
is going to be fine-tuning our weights
in our cross-attention,
as well as in our Vision Transformer in our DistilGPT2.
Let's go ahead and jump right into our code
to see how this is going to work.
In order to fine-tune an image captioning system,
we're going to need a few more imports than normal.
From our Transformer's library,
we'll be encoding a new package or class rather,
called VisionEncoderDecoderModel.
As usual, transformers tends
to have a lot of prebuilt classes for us to use
in order to build what we wanna build.
We'll be also importing a GPT2Tokenizer
and AutoFeatureExtractor, which we've seen previously
in how we use feature extractors,
like tokenizers to process images,
and AutoTokenizer, a TrainingArguments, and our Trainer.
A few other imports here,
we are gonna be using a CountVectorizer
to do some analysis a little later on.
We'll be using Pillow's image class
to load up our images from pathways.
We have a few imports here for visualizations
and reading data.
Our datasets module will be importing Torch and NumPy,
and also a new library called Torchvision.
Torchvision comes with several functions and classes
in order to do computer vision.
We're not gonna be getting too deep into the weeds
about how all of these things work.
However, because we are working with images,
we are going to be using some of the more standard objects
in this library, so let's go ahead and get started here.
We're gonna immediately load up
a randomly initialized VisionEncoderDecoderModel.
The encoder decoder model can take in two models,
one for an encoder, one for a decoder.
Our encoder is going to be a Vision Transformer model,
the base version, trained by Google, using 16 by 16 patches,
on images size 224 by 224 pixels,
pre-trained on the image net in21k dataset
that we've talked about in the last session.
Our decoder is going to be a DistilGPT2.
You can think of this almost like in T5,
where we had an encoder and a decoder,
except the difference here is that while we're loading up
our custom encoder decoder model,
we're able to use two pre-trained models from scratch.
One thing I will note as we load this up is,
every time we load something for sequence classification
or GPT for something,
and even here, VisionEncoderDecoderModel
from encoder decoder pretrained,
which is a special method for loading up two models at once,
it will almost always tell us there are some weights
that are not initialized from the checkpoint.
And I wanna call this out here,
because if we read through all of the weights
that have not been initialized,
or rather they have not been loaded,
but have rather been randomly initialized,
I believe they're, if not all,
then almost all of them, cross-attention weights.
And this makes sense because the Vision Transformer
does not have cross-attention.
It is outputting the key and the value matrices as normal.
It's the GPT2 model, the decoder, which is going to take in
and learn weights for combining that
with its own query matrices through cross-attention.
So these transformers h zero transformers h one, h two,
that relates to the DistilGPT2.
In fact, we've seen those transformer.h.0 before
in our GPT session.
So the main thing that we're going
to be fine-tuning in this model are cross-attention weights,
i.e., how is GPT going to connect speaking
with an image being encoded?
That's the goal.
GPT2 already knows basic language
through its auto aggressive language modeling pre-training.
And Vision Transformer already knows
what images tend to look like through its pre-training
by classifying the ImageNet dataset.
In our fine-tuning, however,
we need to combine these two items through cross-attention.
So we can take a look at the type of encoder.
It's a ViT Vision Transformer model.
Our type of de coder is a GPT2LMHeadModel,
the same type that we have used in the past.
We can add up our total parameters in our entire model.
And we see that our model
has a total of 182 million parameters,
which by today's standards is relatively small.
And that's mostly due to the fact
that we're using a DistilGPT2, not a Baseline GPT2.
We'll instantiate a tokenizer
so that we can tokenize our captions eventually.
And then finally, it's time to load up our images.
So this function, again, is simply to do
with loading up images from our dataset.
Again, the dataset that we're going
to be using is not ImageNet.
However, it is going to be the Flickr8K
dataset from Kaggle.
The dataset has thousands of images,
and each image has a few possible captions
that relate to it.
So this function is simply going to load up those images
and those paths in order to create a Pandas data frame,
where we have an image pathway and a caption.
So this function is gonna output a list of tuples
where the first element is a just a path to an image.
And the second element are multiple,
potentially, multiple possible captions,
that will all be correct for that specific image.
So now that we have this list of pathways
and potential captions, let's go ahead
and load up our featured extractor.
Again, our feature extractor
is like a tokenizer, but for images.
We're going to load the AutoFeatureExtractor
think auto model or AutoTokenizer.
It's going to select the best feature extractor possible
for the same model identifier
that we have for our Vision Transformer.
So like we would load up a tokenizer for bert-base-uncased
to make sure it's the same as our model,
bert-base-uncased, same thing.
We wanna make sure we're loading the same feature extractor
as we are the actual Vision Transformer.
So that's what we're doing here.
This feature extractor is what's going
to transform our raw images from pathways to pixel values,
i.e., tensors of pixel values.
We can take a look at our feature extractor,
has some parameters.
It normalizes the images to the these mean
and standard deviation.
And we also are resizing images to the size 224 by 224.
Now that we have our feature extractor,
we're going to create something called a compose object.
A compose object is going to contain several steps
to pre-process an image.
Now, I know I'm throwing a lot at you
specific to computer vision.
However, just to be clear, I wanna make sure
that we're all aware that we don't have to go too deep
into exactly what each of these items does,
but to make sure that we have a general sense
of what's happening, this compose object
is going to take in a raw image and do three things to it.
First, it's going to randomly crop that image
and resize it down to the feature extractor size,
which is 224.
So it will create a 224 by 224 image.
Next, we're going to convert it to be a tensor.
So, converting it to a tenor will transform it
to be a Torch tensor.
And then normalize it by the mean and standard deviation.
That's basically gonna subtract the mean
and divide by the standard deviation given
for each of the three channels,
because we have R, G, B channels,
each R, G and B value is going to be subtracted by a 0.5
and normalized to that mean and standard deviation.
The RandomizedResizedCrop
to add is a data augmentation tool.
By cropping the image randomly,
what we're actually doing is making the task
of image captioning just a little bit harder,
but it forces the model to generalize a little bit easier.
So, because each image has generally
more than one possible caption,
we will actually end up with technically different images
mapping to similar but different captions,
but we are gonna not necessarily have the risk
of over fitting by having the same image multiple times,
because we're randomly cropping the image
in different places.
So if we iterate over our captions mapping,
we will end up with a Pandas data frame
of nearly 20,000 image caption tuples, basically.
Now that we have our image data frame,
in fact, we could probably just take a quick peek
at what this would look like.
We have a path column and a caption column.
So we see we have multiple of the same pathways
with different captions.
But when we run our pre-processing function over it,
these same pathways will become technically
different cropped images.
So our pre-processing function
is going to be slightly more cumbersome
than our previous ones, in that instead of input IDs,
our Vision Transformer is expecting pixel values.
So to create pixel values,
we are going to call our transforms,
which is our compose object,
over the Pillow image object for that path.
So we're gonna open the path
and then apply our transformations to it
for every element in the batch of paths.
And that's gonna be our pixel values.
Similarly, we're going to use our GPT2 tokenizer
to tokenize our captions.
And we are going to, in this function,
pad them to the maximum length.
And by padding it here, we aren't going
to necessarily need a data collator later on.
Last, we're going to set our labels.
Our labels are going to be technically
the input IDs of our caption.
Now, from a previous session, we have to recall
that because we are padding it,
we need to set the labels to be negative 100,
if it has been encoded into a pad token,
because we don't wanna calculate the loss
based on the pad tokens, because that wouldn't be fair.
That's not actually part of the caption.
So to create our labels here,
we are going to run over our tokenized captions.
And then change any of the labels to be negative zero,
sorry, negative 100, if it is a pad token,
or else, leave it alone.
And then finally, just for posterity's sake,
I will delete the path and caption items
because they're technically not needed anymore.
And I don't wanna have them in memory on my machine.
So when we map our image pre-process function
over our dataset, we end up with a dataset
with only two features in it.
And once we actually run this train test split,
we can see that our dataset has a trainer test
with pixel values and labels.
These are the two items that our model will be expecting.
Remember, we're not using input IDs
because the thing we are inputting
is not text, it's an image.
So the Vision Transformer is expecting pixel values,
not input IDs.
The labels is going to be a list of token IDs
where it's negative 100,
where we don't actually wanna calculate loss,
i.e., the pad tokens.
And also a quick note, I did do a test set of 10%,
as opposed to our normal 20%,
because this is a relatively complex system.
I did wanna skew it a little bit more for training data
than I have for testing data.
We're almost at our training phase.
Last thing that's a little specific to our model here is,
because we have a model that is a combination
of a GPT and a Vision Transformer,
we want to set a few configurations in our model object,
specifically our pad tokens and our decoders start tokens.
We haven't talked about these too much
because they come standard with transformers.
But when you are creating a custom model
that is based on a GPT2 model,
we will want to set the model's pad token
to be the same as the GPT2's pad token.
Similarly, we wanna set the start token
to be the same as the BOS or beginning of sentence token.
Long story short, this is basically how we're going
to let our VisionEncoderDecoderModel
have the same configuration as our GPT2 model,
because when we combine them together,
it does not automatically take the configuration
from the decoder, almost done.
Last thing, before we get
to our actual training arguments down below is,
this is going to be a large and slow training process.
So to speed things up slightly, I'm going to freeze all
but the last two layers in the Vision Transformer.
So to do that, like we've done in the past,
I'll iterate over the named parameters
and freeze all of them
until I see the name encoder layer 10,
because 10 and 11, zero indexed,
are the last two encoders in our Vision Transformer.
You'll note, I'm not freezing anything in the GPT2 model.
And there's a really good reason for that.
And recall, when we instantiated our model,
it had a huge warning that said,
look at all these randomly initialized weights
just so you know you need to fine-tune them.
And those cross-attention weights
are littered all over GPT2.
So I'm not going to freeze a single part of GPT2.
It needs to update a lot more than the Vision Transformer.
So I'm going to freeze most of the Vision Transformer,
kind of taking the assumption that it's already pretty good
at processing images.
I'm only gonna unfreeze the last two encoder layers
of the Vision Transformer, so that I can learn
a little bit nuanced ideas of our images,
but I'm going to leave the GPT2 model completely on
and totally updateable.
So this is our first time in our sessions together,
seeing a situation where we're freezing parts of the encoder
but not freezing parts of the decoder.
Finally, we're at our trainer, where at this point,
hopefully, you're feeling much more comfortable
than all the computer vision stuff
we've been talking about previously.
We have an output directory.
You can see I have a v2 there,
'cause I did have a v1 that I was working on earlier
for my own hobby's sake, and overwrite output directory.
We're gonna go over our dataset four times.
I've upped the batch size to 64,
just to speed things up a little bit more.
We're gonna load the best model at end to be true.
I almost always recommend you use that.
Log level info, we're going to log every 50 steps
to make sure that we are accumulating logs.
And then as usual, we're gonna evaluate
and save after every epoch.
Our trainer will have our model,
our arguments, our train test split.
And then all of a sudden we're ready to go.
So we will run our evaluate first.
We always run evaluate first
to see how we're actually doing.
And we're starting with a loss of 5.1, okay.
So as usual, we're going to hope
that that loss goes down over our four epochs.
And when we run our training system here,
We can see, first thing I'll note is on my laptop,
this took eight and a half hours to run.
So while I was sleeping,
I had this model training on my laptop.
But I wanna call it out, only four epochs took eight hours.
In our T5 session, I believe our training session
took eight hours, but I had 20 epochs running over it.
So this is a large, slow and complex pipeline to train.
And keep in mind, we actually froze
part of the Vision Transformer.
So this is with a frozen portion of the model.
We do see, however, that our validation loss plummeted
on our first epoch, which is always great news,
going from 5.1 to 2.8, and then started to slow down.
This is expected, but because image captioning
is such a broad task, so many things could be in this image,
we do wanna keep in mind
that if we wanted our validation loss to go down even more,
it is more likely that more data will be helpful
rather than more epochs,
because of how quickly the validation loss is slowing down.
So as usual, we'll save our model.
And then we can load it up.
So loading up our VisionEncoderDecoderModel
looks pretty much the same as it would before.
We're gonna load it from our pre-trained directory,
our v2 image captioning.
We can take a look at our model configuration.
Everything seems to be in order, we're not getting an error.
Let's scroll down slightly.
Everything seems to be working.
So now that we have our model, and we're also,
sorry, just a note, we are not getting that warning anymore
where our cross-attentions
are no longer randomly initialized.
We are perfectly loading up everything.
All of the weights were initialized from this model,
which is exactly what we want to see.
So last thing I'm going to do to actually use this system
is I'm going to create a new transform.
The reason I'm doing this is, in our training time,
I had, instead of resize, I had random resize crop,
which made the training a little bit harder
by taking subsets of the image.
When it's time for inference,
I don't wanna take a random crop of that image.
And that's because any random crop
could lead to inconsistencies in our model.
So at inference time, i.e.,
we're not actually training anymore, at inference time,
I'm going to give it a different composition,
our set of composition functions,
that will normalize and convert to a tensor the exact same.
The only difference is we're going to resize
to be that specific 224 by 224, without taking any crops.
And now we can actually use our model.
I wrote a little helper function here
to capture an image, given a model and a path.
That path could be either a URL
in which it will use requests and BytesIO
to convert that into our image.
And we can see that we're running our inference transform
across our image here.
If it's not a URL, it's just a file path,
it will simply open it and apply our image transform.
Unsqueeze 0, as we know, will create a forced accretion
of a batch dimension to the beginning
so that our model can actually take it in.
Once we have our image matrix,
we will actually be able to simply generate,
like we would previously,
when we do model dot generate or pipelines wrapper.
We can do model dot generate,
give it our pixel values, our image matrix.
And then we have the same parameters for decoding,
number of beams, max length, early stopping, sample, top k,
and how many sequences to return.
I'll be honest, off screen,
I played around with some of these parameters,
but I'll give you the intuition
about why each one the way they are.
Number of beams, three.
By default, the number of beams is one,
that greedy search that we saw in our GPT session.
By setting number of beams to be greater than one,
I'm intuitively asking our image caption system
to think ahead a little bit.
Think a few words ahead, and then use that,
take that into consideration.
That generally makes our model a little bit less random
and forces it to think ahead
or what words would I have used in the future
so that I don't necessarily find myself pigeonholed
into a part of the generation tree
that is not going to actually capture the image.
I set a max length of 15 tokens.
I feel like 15 tokens is plenty to describe an image.
Early stopping, I almost always set to true.
If it hits the end of sentence token, just stop trying.
Sampling, I'm setting sample to be true
because I do wanna return multiple possible sequences.
So I'm going to let it choose the next word randomly.
Because I'm letting it choose the next word randomly,
I do wanna sharpen up some of those probabilities
by taking only the top 10 tokens
and then redistributing the softmax,
the distributions across those top 10 tokens.
Lastly, I will decode each of those generated sequences,
display the image and print out our captions.
So all this is going to do is load the image,
run our inference transforms over it,
generate five possible captions,
display our image, print the captions.
So let's go ahead and give it a try.
Our first image here, we see this,
what seems to be a very cool dude.
It's probably the coolest dude I've ever seen,
I dunno about you.
He seems to be standing by a wall with a skateboard.
Very cool, he seems to know how to skateboard,
I believe that.
Now, if we take a look at the possible captions
our model has come up with, it has,
in several variations, a man in a shirt.
It can't seem to decide whether my shirt is red or blue.
It's definitely more blue, it's technically gray.
It's definitely more blue though.
So three outta five times at a blue shirt.
And then the rest of it is about something
to do with skateboarding,
is skateboarding on a sidewalk, okay.
Sidewalk's not technically right, I'm inside, but fine.
Is doing a trick on a skateboard.
I wouldn't really call this a trick,
but you can see the general sense
that the image captioning system
is understanding there's a man in this photo,
he's wearing a shirt, good,
and is doing something with a skateboard.
So all in all, it's not looking too bad.
Now let's try another image.
So we have an image here of a woman
doing yoga on some rocks by some water.
So if we take a look at our captions here,
we can see, fascinating.
A man is rock climbing on a cliff,
overlooking the ocean surface, okay?
Now some things are correct.
She's not technically rock climbing, but okay.
She's on a rock, on a cliff, overlooking the ocean,
not ocean floor, per se.
One of 'em had said desert, near the ocean.
I can see why it might think that, but not true.
I'm a little bit more concerned
with the fact that five outta five times,
it's not saying woman, it's saying man,
which, okay, is interesting,
but it is a good opportunity to reiterate
that as we are fine-tuning these kinds of models,
especially on relatively complex tasks,
like image captioning,
where almost literally there's an infinite world
of possibilities of what could be in the image,
there may be some bias.
And just to really quick check on what's going on,
this code here is going to use the CountVectorizer object
in Scikit-learn to count the number of words
in our captions.
When we do that, our top captions by count can be seen here.
The most common word is dog.
So our images are mostly of dogs.
Then we have man, boy, dogs, and water.
So our top five words in all of these captions
are around dogs, men and water.
Only after our top five do we even see the first girl,
our first color, the word people,
playing, running, and then woman.
So the term man appears 2,800 times.
The word woman only appears 1,000 times.
Similarly, boy appears 1500 times; girl, close, 1300.
Now this is not necessarily meant
to say our dataset is completely biased,
we shouldn't use it.
However, by simply looking
at the count of words in our captions,
I can make a pretty strong case
as to why our image captioning system
is so eager to say man instead of woman,
because most of the data it's seen is about captioning men.
So it's more likely to caption about a man.
Let's take one of the images
from our original flicker dataset.
So hopefully it will do pretty well on this.
And we do see a brown dog
running through a field of brown grass,
which actually is very specific
because it is mostly brown on the bottom here,
next to a pond in the woods.
I don't see a pond, but I could see a world
where maybe this was water in the background.
I don't think it is, but, you know, it's not too bad.
So our image captioning system
may not be state of the art deployed immediately,
start a company around it.
However, with about 30 cells in a Jupyter Notebook,
we are able to see that we can fine-tune
an image captioning system relatively easily,
and with the same ease as we would
a sequence classification model with BERT.
Just a final sanity check,
let's take a non fine-tuned vision encoder model
from Google DistilBERT,
and then use that to caption the same dog image,
just to make sure that it's actually not able
to just do this out of the box.
And our captions that we get with a non fine-tuned version
are definitely not captioning what's in the image.
For some reason, it really wants
to talk about the Department of Homeland Security
and how someone was arrested after shooting someone.
You know, bias is a scary thing.
But the point here is that our model
that has not been fine-tuned
clearly cannot do any image captioning.
So we can definitively say
that our fine-tuning process was successful.
So even though this was our most complicated system
and probably our most complicated NLP task,
we can see how a lot of the patterns of training
and fine-tuning hold true across our different tasks.
Really the only differences here
were how we process the actual image,
which can be reused if you wanted
to try this on an image system at home as well.
So hopefully this gives you a good sense
of how to fine-tune your very own image captioning system.