Building your own object detector — PyTorch vs TensorFlow and how to even get started? | by Maximilian Beckers | Towards Data Science

Having experienced some of the difficulties and headaches of implementing AI systems for object detection, I wanted to share some of the knowledge I gained on how to get started. The first step is always the hardest which is why this very practical approach is designed to ease people into the world of TensorFlow’s as well as PyTorch’s object detection frameworks.
As I already was an experienced data scientist and had developed a few productive machine learning software components in the past years, I thought it should be a fairly easy job to get online, find a pre-trained CNN and train it further on an unseen data set to enable it to detect some previously unknown objects. Up until about a a year ago I had mostly worked with tree-based models, a lot of scikit-learn as well xgboost and of course tons and tons of pandas. As with many AI tasks, my first approach with this one turned out to be a classic version of a: “Not so fast!“ There are a few stepping stones on the way that you have to know of, that however not many articles and blogs seem to mention. After having spent many hours on this topic and having read a lot of TensorFlow source code I know the answers to questions like:
“Why the heck is my example code from the official GitHub repos not working?”
So I thought why not share my experience juggling with these tools with others and describe how I went about solving the issues I faced.
So simply said, my goal was to develop my own object detector — a neural network that can classify and detect objects in an image, not just an off-the-shelf object detection algorithm, but one that would specifically be good at detecting images from my own data set — meaning I would have to do transfer learning.
First of all, I went online, and freshened up my knowledge on how some of these algorithms work mathematically (I can really recommend Andrew Ng’s deeplearning.ai course on Convolutional Neural Nets on coursera. It has a full week dedicated to Object Detection that gets you up and running on the key concepts and nicely describes the math behind everything). So with the theory in mind, I had a look at some of the open-source data sets I could use for my task and stumbled across the following Raccoon data set by experiencor on GitHub https://github.com/experiencor/raccoon_dataset that I then decided to use. It already has annotations as well as images stored away neatly in respective folders and only a single class (raccoon) to detect and therefore makes a nice starting point.
That being done, I had a look at two widely known deep learning frameworks, that let you use pre-trained networks for transfer learning (further training these networks to tailor them to your specific data set), TensorFlow’s object detection API as well as PyTorch’s torchvision package. And one thing I can say right away: I was able to complete my work in PyTorch in about half a day whereas figuring everything I wanted out in TensorFlow took me a few days! The fact that the tf object detection API was built on TensorFlow 1.x and we now already have version 2.x gave me so many compatibility problems that made it practically impossible to get stuff done without extensively going through their source code as some of the examples on their repo were based on 2.x whereas the core training code was still based on 1.x.
I did my coding on Google Colab since they offer you a neat notebook experience and access to free CPU and GPU on their cloud as well as free storage in the form of google drive, that can easily be mounted to your Colab space.
Getting the data was rather easy as I just had to navigate into a different folder (in my mounted google drive storage) and run
! git clone https://github.com/experiencor/raccoon_dataset.git
Inspecting the data then works easily like such:
from PIL import Imageimage = Image.open("<your path>/raccoon_dataset/images/raccoon-1.jpg")image
When it comes to the labels for this image, we have to look inside the annotations folder. In case of object detection, a class for each object (in our case just the one for raccoon) and 4 coordinates per object that represent the bounding box make up the label. To have a quick look you can simply download the raccoon-1.xml file to your laptop and open it up with an editor (I used Atom). You then see the following:
<annotation verified="yes"> <folder>images</folder> <filename>raccoon-1.jpg</filename> <path>/Users/datitran/Desktop/raccoon/images/raccoon-1.jpg</path> <source>  <database>Unknown</database> </source> <size>  <width>650</width>  <height>417</height>  <depth>3</depth> </size> <segmented>0</segmented> <object>  <name>raccoon</name>  <pose>Unspecified</pose>  <truncated>0</truncated>  <difficult>0</difficult>  <bndbox>   <xmin>81</xmin>   <ymin>88</ymin>   <xmax>522</xmax>   <ymax>408</ymax>  </bndbox> </object></annotation>
which tells you that this annotation file corresponds with the raccoon-1.jpg file, the class of the object is “raccoon” and the bounding box edges are represented by xmin, xmax, ymin and ymax. A very natural reaction would be to first check whether this is correct by simply drawing the box into the image.
Pretty easily done by:
from PIL import ImageDrawxmin = 81ymin = 88xmax = 522ymax = 408draw = ImageDraw.Draw(image)draw.rectangle([(xmin, ymin), (xmax, ymax)], outline =”red”)image
So we don’t have to go through every file in the folder and collect the correct labels on each image there’s already a csv file containing all this info. Having a quick look at it can be done like this:
import pandas as pdlabels = pd.read_csv("<your path>/raccoon_dataset/data/raccoon_labels.csv")labels.head()
This is the starting point for both the TensorFlow as well as the PyTorch libraries and also where things are starting to differ between the two.
Section A: TensorFlow
Starting with TensorFlow object detection, is basically supposed to work like this: You clone their repo and install the API according to their installation guide. In practice what you have to do with Google Colab is in some parts easier and in some harder:
Install TensorFlow and all other mentioned libraries (in Google Colab the magic command to pick a tf version is all you need, everything else is preinstalled)
Install TensorFlow and all other mentioned libraries (in Google Colab the magic command to pick a tf version is all you need, everything else is preinstalled)
%tensorflow_version 1.x
2. COCO API installation — is already preinstalled and can be checked with
import pycocotools
3. Clone their repo
! git clone https://github.com/tensorflow/models.git
4. Navigate to the models/research folder
import os os.chdir("<your path>/models/research")
5. Install protobufs
! protoc object_detection/protos/*.proto --python_out=.
6. Now to test your installation simply run:
%%bashexport PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slimpython object_detection/builders/model_builder_test.py
7. Super important: every bash command you do in Colab is like starting a new terminal session so it is very important that, every time you use the CLI for tf’s object detection, you MUST include
export “PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim
at the top of your cell!
Since every interaction with the API is in command line fashion (to train a model you call something like):
%% bash export “PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slimpython”python object_detection/model_main.py \    --pipeline_config_path=${PIPELINE_CONFIG_PATH} \    --model_dir=${MODEL_DIR} \    --num_train_steps=${NUM_TRAIN_STEPS} \    --sample_1_of_n_eval_examples=$SAMPLE_1_OF_N_EVAL_EXAMPLES \    --alsologtostderr
and I wanted to look under the hood of this CLI (plus a notebook environment isn’t really built for command-line executions) I added one key thing into my installation on Colab:
8. Enable the python kernel to find all imports made in the model_main.py file (called in the CLI above).
import syssys.path.append("<your path>/models/research/object_detection")sys.path.append("<your path>/models/research/slim/")
9. Make sure you have numpy==1.17 installed, otherwise the evaluation scripts won’t work.
Voila ready to get started!
Having installed the necessary software and downloaded the training data we can finally think about AI. Which convolutional neural network architecture can we use? Which data set is it trained on? In TensorFlow’s object detection API we can choose from a variety of models available in their detection model zoo (love the name for this by the way :) ) trained on different industry and research standard image datasets. Perhaps the most famous one is the so-called COCO data set (common objects in context) that has many different images with labeled classes, bounding boxes, and even masks available as annotations. The data set is used for many competitions and is, quote “designed to push the state of the art in object detection forward”. Cool thing and pretty much what I was looking for. I picked a network called “Faster R-CNN”. I picked it because it is available in both the TensorFlow model zoo as well as the torchvision package and that way I could compare the two a little easier. Faster R-CNN is a region-proposal network (hence the R) that uses the technique of “anchor boxes” to localize objects and predict them. In a nutshell it’s a combination of two networks, one that detects regions within the image that have a high probability of containing an object, and the other predicting the bounding boxes as well as class probabilities in these regions. There’s plenty of more detailed explanations online. I personally like this one by Hao Gao.
Once you made a choice on the model you first have to download it, meaning you download the network architecture and pre-trained weights. I chose the “faster_rcnn_inception_v2_coco_2018_01_28” model. This is done by running the following command from the models/research folder:
%%bashwget http://download.tensorflow.org/models/object_detection/faster_rcnn_inception_v2_coco_2018_01_28.tar.gztar -xvf faster_rcnn_inception_v2_coco_2018_01_28.tar.gz
You can now find a folder containing all the model files:
import os os.listdir("faster_rcnn_inception_v2_coco_2018_01_28")
which looks like this:
['model.ckpt.index',  'checkpoint',  'pipeline.config',  'model.ckpt.data-00000-of-00001',  'model.ckpt.meta',  'saved_model',  'frozen_inference_graph.pb']
There you can find checkpointing information which lets you reconstruct the model with the pre-trained weights, thus a model that is ready to do transfer learning. You also have a saved_model folder with an exported version of the downloaded model as well as a frozen_inference graph — just as the name proposes a tf graph that you can use for inference. Last but not least you find an example for a pipeline.config file. In this file you can specify exactly how you want your network to be trained, what training data to use, how many object classes you have, and so on. In here we can customize the pre-trained model to our specific needs.
In TensorFlow’s object detection, the estimator API is used to handle the training and evaluation (and some more stuff like checkpointing etc.) process. This API expects the data in a tf.record format which you need to create first.
Luckily this has already been done by the creator of the raccoon data set. Otherwise we could have transformed it ourselves as well:
!python generate_tfrecord.py --csv_input=data/train_labels.csv  --output_path=data/train.record
We have a train.record and a test.record data set which will be used to feed data in batches to the tf.estimator object during training.
In the pipeline.config file we have to adjust 5 areas. The fine_tune_checkpoint, the train_config, the valid_config, the num_classes and the second_stage_post_processing. Everything else can just stay the way it is. The fine_tune_checkpoint is the path to the pre-trained weights contained in the downloaded checkpoint files. All we do here is set it equal to
fine_tune_checkpoint: "<your path >/faster_rcnn_inception_v2_coco_2018_01_28/model.ckpt"
Don’t worry in case you find that there is no file exactly called model.ckpt, trust me it does work this way ;) The model can be rebuilt and the weights are inserted into the correct layers. The num_classes field is set to 1 (1 class only: raccoon).
In the train_config section we can adjust all things concerning the training phase. I left everything at its default and only changed the necessary parts like the train_input_reader, where you have to set the input path to your train.record training data set as well as the label_map_path to a pbtxt file containing info on the labels. I first changed the label_map file to
item {  id: 1  name: 'raccoon'}
then I set the train_input_reader to
train_input_reader: {  tf_record_input_reader {   input_path: "<your path>/raccoon_dataset/data/train.record"  }  label_map_path: "<your path>/training/raccoon_label_map.pbtxt"}
Next I changed the valid config section. Again the name explains it all here. I just set the eval_input_reader to my label_map file as well as the test.record test data set to
eval_input_reader: {  tf_record_input_reader {   input_path: "<your path>/raccoon_dataset/data/test.record"  }label_map_path: "<your path>/training/raccoon_label_map.pbtxt"  shuffle: false  num_readers: 1}
Lastly I adjusted the second_stage_postprocessing. Here you can define what the model does with the output. Generally a detection model outputs a number of possible combinations of class probabilities and bounding box coordinates (e.g. one set per anchor) that is not equal to the number of objects in the image. As users we then have to use techniques such as non-max suppression to get the best-fit bounding boxes (Andrew Ngs Coursera course as mentioned above gives a really nice introduction to this!) to every object in the image. I picked the following here:
second_stage_post_processing {      batch_non_max_suppression {        score_threshold: 0.5        iou_threshold: 0.6        max_detections_per_class: 5        max_total_detections: 5      }      score_converter: SOFTMAX    }
meaning I will only allow 5 detections per image (ideally 5 different detections of different raccoons if existent in the image). The IoU (intersection over union) threshold, as well as the score threshold, are tailor towards only having a single bounding box per object.
Once the pipeline.config file is good to go we can start the training. TensorFlow suggests doing this with a command-line job calling their model_main.py file with a variety of input parameters. However, as I wanted to look under the hood a little bit, I hopped into this file and ran the commands myself. For the training and evaluation of the model the object detection API is making use of tf’s estimator API. It takes care of training and evaluation for us nicely and integrates automatic checkpointing as well as model exporting. It also allows for transfer learning from a given checkpoint (which is exactly what we are doing) and writes automatic summaries during training that we can visualize with tensorboard to monitor performance.
The estimator will need training and evaluation specs that incorporate and define all the things just mentioned and takes care of batching the data for us. tf’s object detection has predefined functions to instantiate these objects from our pipline.config file:
Before we can start we need to import all the necessary libraries (just as it is done in the model_main.py file):
from __future__ import absolute_importfrom __future__ import divisionfrom __future__ import print_functionfrom absl import flagsimport tensorflow as tffrom object_detection import model_hparamsfrom object_detection import model_lib
Now we can start by specifying a model directory:
model_dir = "<your path>/faster_rcnn_inception_v2_coco_2018_01_28/model_dir"
The API will write its log files as well as checkpoints and exported models to this directory. Then after specifying a run config, the path to our pipeline.config file, as well as the number of steps (one step is one run through a single batch of data) to train for:
config = tf.estimator.RunConfig(model_dir=model_dir)pipeline_config_path= "<your path>/training/pipeline.config"num_train_steps=10000
we can instantiate all the objects and finally create the train and eval specs using the appropriate pre-built function.
train_and_eval_dict = model_lib.create_estimator_and_inputs(                        run_config=config,                        hparams=model_hparams.create_hparams(None),                        pipeline_config_path = pipeline_config_path,                        train_steps =num_train_steps,                        sample_1_of_n_eval_examples = 1)estimator = train_and_eval_dict['estimator']train_input_fn = train_and_eval_dict['train_input_fn']eval_input_fns=train_and_eval_dict['eval_input_fns']eval_on_train_input_fn=train_and_eval_dict['eval_on_train_input_fn']predict_input_fn = train_and_eval_dict['predict_input_fn']train_steps = train_and_eval_dict['train_steps']
Estimator, train_spec and eval_spec are then used for training:
tf.estimator.train_and_evaluate(estimator,train_spec,eval_specs[0])
A few things happen now that are worth mentioning. First of all the model gets trained using the appropriate optimizer settings from our pipeline.config file. Every config.save_checkpoints_secs the model writes checkpoints (a set of weights) to the model_dir. Then every eval_specs[0].throttle_secs we pause the training to run evaluation on the test set. In order for this to work we need to have a checkpoint available first. The eval_spec is configured to use the evaluation metrics of the COCO challenge (using the pycocotools library). Every config.save_summary_steps the estimator API writes out event files for visualization in tensorboard.
I ran the model for 10.000 steps and got a final mAP of 0.65 on the test set. This already outputs some pretty decent bounding boxes when making predictions for the naked eye. But in order to see that, we first need to figure out how to do inference with this trained model.
This is perhaps the hardest part to figure out in the whole object detection pipeline. In TensorFlow’s object detection repo there are some examples on how to do inference on pre-built models, however, the code relies on TensorFlow version 2.x. Using a saved model or a frozen inference Graph with TensorFlow 1.x code is (relative to tf 2.x) a lot more complicated since you have to work directly with the tf graph and session.
Either way we first have to export the model in the correct fashion. It took me some experimenting to find out why the exporter that was already in the eval_specs and exports a saved model at the end of tf.estimator.train_and_evaluate doesn’t do the job. Essentially this is due to the serving function used in this automatic exporter. It is tailored towards using tf.example type data (e.g. test.record data), however, we are trying to input a 3D tensor representing an image. This is why we have to run an export job where we specify the input type.
With the following imports:
import tensorflow as tffrom google.protobuf import text_formatfrom object_detection import exporterfrom object_detection.protos import pipeline_pb2slim = tf.contrib.slimflags = tf.app.flags
we have to provide an output directory, our pipeline.config file, the input type as well as a checkpoint from which we would like to export the model (just use the latest one)
pipeline_config_path = "<your path>/training/pipeline.config"config_override = ""input_shape = Falsetrained_checkpoint_prefix ="<your path>/faster_rcnn_inception_v2_coco_2018_01_28/model_dir/model.ckpt-10000"output_directory = "<your path>faster_rcnn_inception_v2_coco_2018_01_28/output"input_type = "image_tensor"write_inference_graph = False
and now we can run the exporter
tf.reset_default_graph()pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()with tf.gfile.GFile(pipeline_config_path, 'r') as f:   text_format.Merge(f.read(), pipeline_config)text_format.Merge(config_override, pipeline_config)if input_shape:   input_shape = [        int(dim) if dim != '-1' else None        for dim in input_shape.split(',')   ]else:    input_shape = Noneexporter.export_inference_graph(   input_type, pipeline_config, trained_checkpoint_prefix,   output_directory, input_shape=input_shape,    write_inference_graph=write_inference_graph)
For TensorFlow 2.x all we have to do is reload the runtime and select 2.x as version
%tensorflow_version 2.x
and pretty much just follow along their example notebook. You can then load the model with two lines of code:
model = tf.saved_model.load("<your path>/faster_rcnn_inception_v2_coco_2018_01_28/output/saved_model")model = model.signatures['serving_default']
All we have to do now is load up an image from the test set (e.g raccoon-14.jpg)
from PIL import Image, ImageDrawimport numpy as npimport pandas  as pdfilename = "raccoon-14.jpg"img = Image.open("<your path>/raccoon_dataset/images/"+filename)test = pd.read_csv("<your path>/raccoon_dataset/data/test_labels.csv")true_labels = test[test["filename"] == filename][["xmin", "ymin", "xmax","ymax"]]
and run the inference (as suggested by the tf inference scripts)
image = np.asarray(img)# The input needs to be a tensor, convert it using `tf.convert_to_tensor`.input_tensor = tf.convert_to_tensor(image)# The model expects a batch of images, so add an axis with `tf.newaxis`.input_tensor = input_tensor[tf.newaxis,...]# Run inferenceoutput_dict = model(input_tensor)# All outputs are batches tensors.# Convert to numpy arrays, and take index [0] to remove the batchdimension.# We're only interested in the first num_detections.num_detections = int(output_dict.pop('num_detections'))output_dict = {key:value[0, :num_detections].numpy()                 for key,value in output_dict.items()}output_dict['num_detections'] = num_detections# detection_classes should be ints.output_dict['detection_classes'] = output_dict['detection_classes'].astype(np.int64)
Now we can draw the learned bounding box(es) in red and the true labels in green:
selected_boxes = output_dict["detection_boxes"]selected_scores = output_dict["detection_scores"]draw = ImageDraw.Draw(img)#draw predictions:for i in range(len(selected_scores)):  draw.rectangle([(selected_boxes[i][1]*img.width,                     selected_boxes[i][0]*img.height), (selected_boxes[i][3]*img.width,  selected_boxes[i][2]*img.height)],  outline ="red",  width = 3)  draw.text((selected_boxes[i][1]*img.width, selected_boxes[i][0]*img.height), text = str(selected_scores[i]))# draw groundtruth:for i in range(len(true_labels["xmin"].values)):  draw.rectangle([(true_labels["xmin"].values[i], true_labels["ymin"].values[i]), (true_labels["xmax"].values[i], true_labels["ymax"].values[i])],  outline ="green",  width = 3)img
and we can see a pretty cool little raccoon doing some major housekeeping. Also the bounding boxes and probabilities are pretty good considering we only trained the network for 10.000 steps.
Note: If you want to do this with tf version 1.x you can follow the scripts located here.
Section B: PyTorch
PyTorch has a package called torchvision that includes model architectures, data sets, and other helpful functions for computer vision. Torchvision also has a subpackage on object detection which we will be using in this section. A lot of the following setup and code is modeled according to torchvision’s object detection tutorial.
We need to make sure that we have numpy version 1.17 installed.Check pycocotools installation
We need to make sure that we have numpy version 1.17 installed.
Check pycocotools installation
import pycocotools
4. Change to a PyTorch object detection directory you created
import osos.chdir("<your path>/pytorch object detection")
3. Clone the PyTorch vision repo and copy some python files
%%bashgit clone https://github.com/pytorch/vision.gitcd visiongit checkout v0.3.0cp references/detection/utils.py ../cp references/detection/transforms.py ../cp references/detection/coco_eval.py ../cp references/detection/engine.py ../cp references/detection/coco_utils.py ../
4. Now we can do all the necessary imports
import numpy as npimport torchimport torch.utils.datafrom PIL import Imageimport pandas as pdfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictorfrom engine import train_one_epoch, evaluateimport utilsimport transforms as T
As opposed to having to create a new type of data file for our training data (as with the training.record file in TensorFlow) we only have to write a data set python class that is later used by the model to parse the images and corresponding annotations. The class needs to inherit from torch.utils.data.Dataset and implement the __getitem__ and __len__ methods. As we are just going to use the same raccoon data set as earlier and the data folder in the repo already contains a raccoon_labels.csv file containing all the parsed annotations, all we have to do is write a little helper function:
def parse_one_annot(path_to_data_file, filename):   data = pd.read_csv(path_to_data_file)   boxes_array = data[data["filename"] == filename][["xmin", "ymin",           "xmax", "ymax"]].values      return boxes_array
which we can then use to write our RaccoonDataset class:
class RaccoonDataset(torch.utils.data.Dataset):   def __init__(self, root, data_file, transforms=None):       self.root = root       self.transforms = transforms       self.imgs = sorted(os.listdir(os.path.join(root, "images")))       self.path_to_data_file = data_filedef __getitem__(self, idx):      # load images and bounding boxes      img_path = os.path.join(self.root, "images", self.imgs[idx])      img = Image.open(img_path).convert("RGB")      box_list = parse_one_annot(self.path_to_data_file,       self.imgs[idx])      boxes = torch.as_tensor(box_list, dtype=torch.float32)num_objs = len(box_list)      # there is only one class      labels = torch.ones((num_objs,), dtype=torch.int64)      image_id = torch.tensor([idx])      area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:,      0])      # suppose all instances are not crowd      iscrowd = torch.zeros((num_objs,), dtype=torch.int64)      target = {}      target["boxes"] = boxes      target["labels"] = labels      target["image_id"] = image_id      target["area"] = area      target["iscrowd"] = iscrowdif self.transforms is not None:         img, target = self.transforms(img, target)      return img, targetdef __len__(self):         return len(self.imgs)
We can run a quick check to see if everything is implemented correctly. Instantiate a RaccoonDataset object and run __getitem__
dataset = RaccoonDataset(root= "<your path>/raccoon_dataset",data_file= "<your path>/raccoon_dataset/data/raccoon_labels.csv")dataset.__getitem__(0)
and the output should look like this
(<PIL.Image.Image image mode=RGB size=650x417 at 0x7FE20F358908>,  {'area': tensor([141120.]), 'boxes': tensor([[ 81.,  88., 522., 408.]]),    'image_id': tensor([0]),    'iscrowd': tensor([0]),    'labels': tensor([1])})
Since we used a faster R-CNN earlier for TensorFlow we will also use one here. Also it’s not like torchvision gives us a big choice. They only have 2 pre-trained (on coco) models for object detection available on their website. A Faster R-CNN ResNet-50 FPN and a Mask R-CNN ResNet-50 FPN.
In order to download the pre-trained faster r-cnn we can write a little get_model function that adjusts the last layer of the model to output a number of classes that suits our raccoon data set:
def get_model(num_classes):   # load an object detection model pre-trained on COCO   model = torchvision.models.detection.           fasterrcnn_resnet50_fpn(pretrained=True)# get the number of input features for the classifier   in_features = model.roi_heads.box_predictor.cls_score.in_features   # replace the pre-trained head with a new on   model.roi_heads.box_predictor = FastRCNNPredictor(in_features,/   num_classes)      return model
We can also add a little transformer function that enhances our training set by doing some basic data augmentation (horizontal flip of images)
def get_transform(train):   transforms = []   # converts the image, a PIL image, into a PyTorch Tensor   transforms.append(T.ToTensor())   if train:      # during training, randomly flip the training images      # and ground-truth for data augmentation      transforms.append(T.RandomHorizontalFlip(0.5))   return T.Compose(transforms)
Now we can instantiate our training and testing data classes and assign them to data_loaders that control how images are loaded during training and testing (batch size etc).
# use our dataset and defined transformationsdataset = RaccoonDataset(root= "<your path>/raccoon_dataset",          data_file= "<your           path>/raccoon_dataset/data/raccoon_labels.csv",          transforms = get_transform(train=True))dataset_test = RaccoonDataset(root= "<your path>/raccoon_dataset"               data_file= "<your               path>/raccoon_dataset/data/raccoon_labels.csv",               transforms = get_transform(train=False))# split the dataset in train and test settorch.manual_seed(1)indices = torch.randperm(len(dataset)).tolist()dataset = torch.utils.data.Subset(dataset, indices[:-40])dataset_test = torch.utils.data.Subset(dataset_test, indices[-40:])# define training and validation data loadersdata_loader = torch.utils.data.DataLoader(              dataset, batch_size=2, shuffle=True, num_workers=4,              collate_fn=utils.collate_fn)data_loader_test = torch.utils.data.DataLoader(         dataset_test, batch_size=1, shuffle=False, num_workers=4,         collate_fn=utils.collate_fn)print("We have: {} examples, {} are training and {} testing".format(len(indices), len(dataset), len(dataset_test)))
First, we can use the following command to check whether the correct GPU settings are in place:
torch.cuda.is_available()
If you are running on Google Colab just switch the runtime time to GPU.
Now we can set up the model using our get_model function. Also we have to specify an optimizer and a learning rate scheduler, that adjusts the learning rate over time.
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')# our dataset has two classes only - raccoon and not racoonnum_classes = 2# get the model using our helper functionmodel = get_model(num_classes)# move model to the right devicemodel.to(device)# construct an optimizerparams = [p for p in model.parameters() if p.requires_grad]optimizer = torch.optim.SGD(params, lr=0.005,                            momentum=0.9, weight_decay=0.0005)# and a learning rate scheduler which decreases the learning rate by # 10x every 3 epochslr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,                                               step_size=3,                                               gamma=0.1)
In order to do the training we now must write our own for loop over the number of epochs we wish to train on, then call PyTorch’s train_one_epoch function, adjust the learning rate and finally evaluate once per epoch.
# let's train it for 10 epochsnum_epochs = 10for epoch in range(num_epochs):   # train for one epoch, printing every 10 iterations   train_one_epoch(model, optimizer, data_loader, device, epoch,                   print_freq=10)# update the learning rate   lr_scheduler.step()   # evaluate on the test dataset   evaluate(model, data_loader_test, device=device)
In order to save the trained model (there is no automatic checkpointing) we can run the following command once we are done training:
os.mkdir("<your path>/pytorch object detection/raccoon/")torch.save(model.state_dict(), "<your path>/pytorch object detection/raccoon/model")
After 10 epochs I got a mAP of 0.66.
We first have to load the model, which works by instantiating the same “empty” model with our get_model function from earlier and then loading the saved weights into the network architecture.
loaded_model = get_model(num_classes = 2)loaded_model.load_state_dict(torch.load(“<your path>/pytorch object detection/raccoon/model”))
Now we can use our dataset_test class to get test images and their ground truth and run them through the predictor.
idx = 0img, _ = dataset_test[idx]label_boxes = np.array(dataset_test[idx][1]["boxes"])#put the model in evaluation modeloaded_model.eval()with torch.no_grad():   prediction = loaded_model([img])image = Image.fromarray(img.mul(255).permute(1, 2,0).byte().numpy())draw = ImageDraw.Draw(image)# draw groundtruthfor elem in range(len(label_boxes)):   draw.rectangle([(label_boxes[elem][0], label_boxes[elem][1]),   (label_boxes[elem][2], label_boxes[elem][3])],    outline ="green", width =3)for element in range(len(prediction[0]["boxes"])):   boxes = prediction[0]["boxes"][element].cpu().numpy()   score = np.round(prediction[0]["scores"][element].cpu().numpy(),                    decimals= 4)   if score > 0.8:      draw.rectangle([(boxes[0], boxes[1]), (boxes[2], boxes[3])],       outline ="red", width =3)      draw.text((boxes[0], boxes[1]), text = str(score))image
All in all, it is safe to say that for people that are used to imperative style coding (code gets executed when written) and have been working with scikit-learn type ML frameworks a lot, PyTorch is most likely going to be easier for them to start with (this might also change once TensorFlow upgrades the object detection API to tf version 2.x). In this article I am only focusing on the “ease of use” side of things, not looking at performance speed, variability, or portability of the model. A lot of these features however are needed when trying to deploy such models into productive software environments. Here the higher-level API used in TensorFlow with the config.py file incorporates a lot more flexibility from the get-go within the framework, where the lower-level torchvision is held more open, meaning it is easier to look “behind the curtain” (you don’t have to work through a ton of TensorFlow source code), however only the standards are displayed in their tutorial. Anything outside of the standard must be written from scratch, whereas TensorFlow might have a section in the config file to address it. Also the TensorFlow model zoo just has way more pre-trained models available for people to use. Overall my experience resembles a very typical sentiment in the data science community, which is that PyTorche’s object detection framework might be better for research purposes and “understanding every detail” of an implementation, while TensorFlow’s API is a better choice for building productive models.
I hope you guys got some valuable information out of this. Now it’s time to try it out yourselves. Happy coding :)
TensorFlow Object detection:
github.com
Torchvision:
pytorch.org
Convolutional Neural Networks (coursera):