Train a Production Level Image Classifier

Maria Sargsyan, Wolfram Machine Learning Team
Jofre Espigule-Pons, Wolfram Machine Learning Team

(RowDefault[, (Spacer1), , (Spacer1), ])  #WolframTechConf

ABSTRACT

In this presentation, we are going to make a production level image classifier using millions of images. We are going to use Wolfram scripts to train and analyze the trained deep learning models and we will explore several techniques to improve the models.

Introduction

Goals:

Training a net model with tons of data might be unfeasible using a personal computer. Here we will focus on the specific task to improve our current ImageIdentify model.

Improve our current Wolfram ImageIdentify Network Model

The current model was trained back in 2017 using an Inception inspired Neural Net.

NetModel["Wolfram ImageIdentify Net V1"]

Since then there have been several breakthroughs in the research of net architectures. Our primary goal is to achieve a good balance among classification accuracy, size and evaluation speed :

PerformanceGoal -> Speed:Train a model which is both lighter and faster than Wolfram ImageIdentify v1 but yields similar or higher classification results )

PerformanceGoal -> Quality: Train a model which has higher classification results than Wolfram ImageIdentify v1 having similar model size.

Recap into Wolfram ImageIdentify - v1

The neural net used by ImageIdentify can be accessed using NetModel :

NetModel["Wolfram ImageIdentify Net V1"][]

ImageIdentify[]

Get 10 different identifications, along with their probabilities .  Notice that the total of the top results is greater than 1 :

data = ImageIdentify[, Entity["Word", "car"], 10,   "Probability"]

TRAINING SET INFORMATION

Internal Wolfram ImageIdentify training set, consisting of over 3 million training images and over 4, 000 classes of objects (not publicly available) .

Get the label length of Wolfram ImageIdentify v1 :

Length[NetExtract[NetModel["Wolfram ImageIdentify Net V1"],{"Output","Labels"}]]

ImageNet Competition Data and Wolfram ImageIdentify share 530 classes:

Length@Intersection[NetExtract[NetModel["EfficientNet Trained on ImageNet"],{"Output","Labels"}],NetExtract[NetModel["Wolfram ImageIdentify Net V1"],{"Output","Labels"}]]

Wolfram ImageIdentify covers a diverse set of classes:

Dataset[Association@@@Transpose@{Thread["Input Image"->#],Thread["Wolfram Image Identify-v1 outputs"->NetModel["Wolfram ImageIdentify Net V1"][#]],Thread["EfficientNet-v1 outputs"->NetModel["EfficientNet Trained on ImageNet"][#]]}&@{,,,,,,,}]

(Legended[, (SwatchLegend[(EntityTraditionalFramed["artifact", Entity["Concept","Artefact::vnc93"], "Entity[\"Concept\", \"Artefact::vnc93\"]"]), (EntityTraditionalFramed["animal", Entity["Concept","Animal::6w7hz"], "Entity[\"Concept\", \"Animal::6w7hz\"]"]), (EntityTraditionalFramed["flora", Entity["Concept","PlantLife::9y2mf"], "Entity[\"Concept\", \"PlantLife::9y2mf\"]"]), (EntityTraditionalFramed["matter", Entity["Concept","Matter::64zy8"], "Entity[\"Concept\", \"Matter::64zy8\"]"]), (EntityTraditionalFramed["natural object", Entity["Concept","NaturalObject::75p74"], "Entity[\"Concept\", \"NaturalObject::75p74\"]"]), (EntityTraditionalFramed["abstract entity", Entity["Concept","AbstractEntity::4t243"], "Entity[\"Concept\", \"AbstractEntity::4t243\"]"]), (EntityTraditionalFramed["location", Entity["Concept","Location::2h739"], "Entity[\"Concept\", \"Location::2h739\"]"]), (EntityTraditionalFramed["thing", Entity["Concept","Thing::w9d5z"], "Entity[\"Concept\", \"Thing::w9d5z\"]"]), (EntityTraditionalFramed["physical process", Entity["Concept","PhysicalProcess::p82jc"], "Entity[\"Concept\", \"PhysicalProcess::p82jc\"]"]), (EntityTraditionalFramed["formation", Entity["Concept","GeologicalFormation::mth2b"], "Entity[\"Concept\", \"GeologicalFormation::mth2b\"]"]), (EntityTraditionalFramed["causal agency", Entity["Concept","CausalAgency::rj89k"], "Entity[\"Concept\", \"CausalAgency::rj89k\"]"]), (EntityTraditionalFramed["part", Entity["Concept","Part::7726r"], "Entity[\"Concept\", \"Part::7726r\"]"]), (EntityTraditionalFramed["cell", Entity["Concept","Cell::n2k6x"], "Entity[\"Concept\", \"Cell::n2k6x\"]"]), (EntityTraditionalFramed["dry land", Entity["Concept","DryLand::7497m"], "Entity[\"Concept\", \"DryLand::7497m\"]"]), (EntityTraditionalFramed["micro-organism", Entity["Concept","Micro-organism::8vr37"], "Entity[\"Concept\", \"Micro-organism::8vr37\"]"]), (EntityTraditionalFramed["fungus", Entity["Concept","Fungus::2sqjc"], "Entity[\"Concept\", \"Fungus::2sqjc\"]"]), (EntityTraditionalFramed["ground", Entity["Concept","Ground::43s85"], "Entity[\"Concept\", \"Ground::43s85\"]"]), (EntityTraditionalFramed["ice", Entity["Concept","Ice::56ry3"], "Entity[\"Concept\", \"Ice::56ry3\"]"]), (EntityTraditionalFramed["floater", Entity["Concept","Floater::3747w"], "Entity[\"Concept\", \"Floater::3747w\"]"])])])

Choosing a Net Architecture for Wolfram ImageIdentify-v2

Circular Diagram of the Wolfram Neural Net Repository(More Details in “https://community.wolfram.com/groups/-/m/t/1851854”)

ImageNet top1 accuracy vs the number of parameters

Latency measured in seconds on 3.1 GHz Quad-Core Intel Core i7

Training Workflow

The basic Workflow

Running Wolfram Language from the command line

$wolframscript -file  train.wls path_to_config.m

Configuration file

<|
(*------------- Directory/Storage Configurations*)
   (*The Directory under which all the Training Experiment files are stored*)
  "Experiment Root Directory" -> "Path_to_Experiment Directory",
  (*Git Branch Directory*)
  "Git Path" -> "Path_to_Git_Directory",
  "Train Dataset File" -> File[ "..trainData.mx"],
(*------------- Neural Net Configurations
Either arguments for NetModel or filepath to .wlnet object
*)
  "NetModel arguments" -> "ShuffleNet-V2 Trained on ImageNet Competition Data",
  (* In case of Identity, leaves the encoder attached to the net,
    else pass arguments for a new  NetEncoder
  *)
  "NetEncoder Arguments" -> Identity,
  "ExperimentName" -> "ShuffleNet_LabelSmoothing",
  (*"Net Backbone Extraction Function" Either a pure function or Identity*)
  "Net Backbone Extraction Function" ->  (NetTake[#, {1,-3}]&),
  "Add Dropout Before The Prediction Layer" -> <|
    "add or overwrite if exists" -> True,
    "rate" -> 0.2
  |>,
  "Prepend To Net ImageAugmentationLayer" -> False,
  (*Use smooth labels instead of  the hard ones y_new=(1-alpha)*y+alpha/num_classes*)
  "Label smoothing" -> 0.,
(*------------- Training Configurations*)
  "Random Seed" -> 25,
  (* Train Validation Split Ratio should be between 0 and 1*)
  "Train Validation Split Ratio" -> 0.9,
  "NetTrain Params" -> ({
    WorkingPrecision -> "Mixed",
    TrainingProgressReporting -> {{"Print", "Interval" -> Quantity[10, "Seconds"]}, {File[#ProgressFile], "Interval" -> Quantity[5, "Minutes"]}},
    TrainingProgressCheckpointing ->
    {"Directory", #CheckpointDir, "Interval" -> Quantity[1/4,"Hours"]},
    MaxTrainingRounds -> 12,
    TrainingStoppingCriterion -> <|"Criterion" -> "Loss", "Patience"-> 3|>,
    BatchSize -> 64,
    TargetDevice -> {"GPU", 2},
    (* for freezing everything until  the  last layer*)
    Method -> {"ADAM", "L2Regularization" -> 0.1^5},
    LearningRate->0.0001
  }&)
|>

Techniques to improve the training

Image Augmentation

Use ImageAugmentationLayer in combination with PaddingLayer to augment training data during the train time:

ia=NetChain[{PaddingLayer[{{0,0},{4,4},{4,4}},Padding->"Reflected"],ImageAugmentationLayer[{224,224},"ReflectionProbabilities"->{0,0.5}]},
"Input"-> NetEncoder[{"Image",{224,224}}],"Output"->NetDecoder["Image"]]

Define Test Image:

testImg=

Test the Image Augmentation Layer in the Train Mode:

Table[ia[testImg,NetEvaluationMode->"Train"],10]

Notice that the Image Augmentation Layer in the Test Mode does not perform any random Augmentations:

Table[ia[testImg,NetEvaluationMode->"Test"],10]

Attach the ImageAugmentLayer to your net:

NetPrepend[NetModel["EfficientNet Trained on ImageNet"],"ImageAugment"->ia]

Dropout

Some of the nets in the repository do not contain the Dropout layers:

NetModel["ShuffleNet-V2 Trained on ImageNet Competition Data"]

Information[NetModel["ShuffleNet-V2 Trained on ImageNet Competition Data"],"LayerTypeCounts"]

Insert a Dropout layer before the final prediction Layer

NetInsert[NetModel["ShuffleNet-V2 Trained on ImageNet Competition Data"],"top_dropout"->DropoutLayer[0.2],-3]

Label Smoothing

Instead of using hard labels for CrossEntropyLossLayer use Smooth Labels:

crossEntropyWithSmoothLabels[alpha_,labels_]:=Block[
{k=Length[labels],enc=NetEncoder[{"Class",labels}]},
NetGraph[{"embed"->UnitVectorLayer[],"smoothLabels"->ElementwiseLayer[(1-alpha)#+alpha/k&],
"loss"->CrossEntropyLossLayer["Probabilities"]
},{NetPort["Target"]->"embed"->"smoothLabels",{NetPort["Input"],"smoothLabels"}->"loss"},"Target"->enc]
]

ceMnist=crossEntropyWithSmoothLabels[0.1,Range[0,9]]

Prepare MNIST Dataset

mnist=ResourceObject["MNIST"];
trainSet=ResourceData[mnist,"TrainingData"];
valSet=ResourceData[mnist,"TestData"];
report={"Entropy",<|"Measurement"->("ErrorRate"->5),"Key"->"Top5"|>,<|"Measurement"->("ErrorRate"->3),"Key"->"Top3"|>,
<|"Measurement"->("ErrorRate"->1),"Key"->"Top1"|>};

Train MNIST with Default CrossEntropyLossLayer

NetTrain["LeNet",trainSet,All,ValidationSet->valSet,TrainingProgressMeasurements->report]

Train MNIST with the smooth labels. Note that top1-top3 error rates are much smaller:

NetTrain["LeNet",trainSet,All,ValidationSet->valSet,LossFunction->ceMnist,TrainingProgressMeasurements->report]

Transfer Learning

Instead  of training with random weights, retrain the models coming from the ImageNet Competition Data.

It is possible to freeze all the layer except the last using LearningRateMultipliters.

Tutorial: Perform Net Surgery for Transfer Learning

Learning Rate Schedule

Use the learning rate Schedule to dump the learning rate of the model over the epochs:

With "LearningRateSchedule" -> f, the learning rate for a given batch will be calculated as initial*f[batch, total], where batch is the current batch number, total is the total number of batches that will be visited during training, and initial is the initial learning rate specified using the LearningRate option . The value returned by f should be a number between 0 and 1.

Blind Training

Once the model was fully train, retrain the model using the whole training dataset including the validation data.

Notebook Templating for Test Analysis Reports

Video Tutorial: Notebook Templating: An Introduction to Report Generation

Create a Template Notebook with CreateNotebook:

CreateNotebook["Template"]

After creating and editing a custom template we can specify the arguments for a certain trained net:

args=<|"Experiment Name"->"ShuffleNet-v2 2022-09-19",
"Architecture"->"ShuffleNet-v2",
"Probabilities"->"ImageIdentifyNet_Probabilities_NumericArrays_2022-09-19.wxf",
"ResultsObject"->"ImageIdentifyNet_ResultsObject_2022-09-19.wxf",
"LogFile"-> "TrainingLog_2022-09-19T16:59:48.CSV",
"ModelSize"->"28MB",
"ModelSpeed"->"0.3s",
"ImageSize"->"150x100"|>;

Finally we can pass the arguments and generate the report using TemplateApply:

TemplateApply[NotebookTemplate["/Users/jofre/WRI related/ML/WL related/Image Identify/Report-Template.nb"],args]

Analyzing Net Models Validation Loss

Testing the networks on other applications

One-shot learning with Classify

c=Classify[{->"Argiope-bruennichi",->"Argiope-bruennichi",->"Argiope-lobata",->"Argiope-lobata"},FeatureExtractor->NetTake[Import["/Users/jofre/trainedNet.wlnet"],"pool"]]
(CopyTag[ClassifierFunction[]])

c[]

FindClusters

FindClusters[{,,,,,},FeatureExtractor->NetTake[Import["/Users/jofre/trainedNet.wlnet"],"pool"]]

FeatureSpacePlot

FeatureSpacePlot[{,,,,,},FeatureExtractor->NetTake[Import["/Users/jofre/trainedNet.wlnet"],"pool"]]

FeatureSpacePlot[{,,,,,},FeatureExtractor->NetTake[NetModel["Wolfram ImageIdentify Net V1"],"global_pool"]]

FeatureSpacePlot[{,,,,,},FeatureExtractor->NetTake[NetModel["EfficientNet Trained on ImageNet with NoisyStudent"],"avg_pool"]]

Other Resources and future directions

Apply graph optimization to make final NetGraph faster and lighter

Fusing batchnorms with convs

Removing dropouts

Updates on the ontology

Data Curation (Removing placeholder images,...)

Remote Training using RemoteBatchSubmit

Remote Batch Jobs Guide

Training on a AWS server

Image datasets

CIFAR-100 computer-vision training dataset

iNaturalist Open Data

Introduction to Machine Learning

Online book