Skip to content
This repository has been archived by the owner on Dec 10, 2023. It is now read-only.
/ tf-retrain-model Public archive

Retrain a Pretrained Model for New Categories with TensorFlow

License

Notifications You must be signed in to change notification settings

sbrodehl/tf-retrain-model

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

22 Commits
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Retrain a Pretrained Model for New Categories

The full, detailed version of the tutorial(s) can be found on the TensorFlow website

Prerequisites

We will use the flower dataset to retrain the final layer of the inception model. The dataset can be downloaded here. The images are already split up into folders, corresponding to the different classes.

This is for the lazy.

wget http://download.tensorflow.org/example_images/flower_photos.tgz
tar -xf flower_photos.tgz

The pretrained model will be downloaded once at the beginning of the training.

Training

After preparing the data in the corresponding folders one can retrain the final layer of the model, also called transfer learning. See python3 retrain.py --help for more information on training parameters (etc.).

python3 retrain.py --image_dir flower_photos

This will train for a fixed amount of steps. The progress can be seen in tensorboad.

tensorboard --logdir=/tmp/retrain_logs

Visit the local tensorboard website to see pretty things like accuracy or cross entropy graphs.

Inference

Once the retraining of the final layer(s) is done one can use the model to do inference, that is, predict the class for unseen data. Inference is run on a single image (.png/.jpg/.bmp/.gif).

python3 label_image.py --graph=/tmp/output_graph.pb --labels=/tmp/output_labels.txt --output_layer=final_result:0 --input_layer=Mul:0 --image=/FULL/PATH/TO/IMAGE.jpg

Using flower_photos/daisy/100080576_f52e8ee070_n.jpg the output will look like

daisy 0.938932
sunflowers 0.049032
dandelion 0.00581826

Releases

No releases published

Packages

No packages published

Languages