akrantz/powercell-detection repository overview

⁠NOTE: this image is currently broken! DO NOT USE IT!

I am in the process of changing where the dataset is stored and as such, every version of the image is broken.

⁠FRC 2020 Power Cell Detection

This image re-trains the SSD MobileNet v1 or v2 model to detect power cells on the field. By default, it trains MobileNet v1, but you can retrain MobileNet v2 by setting the environment variable NETWORK_TYPE to mobilenet_v2_ssd. Tensorboard is automatically started on startup and can be accessed at port 6006. NOTE: the container currently does not handle CTRL+C gracefully, so you will need to use docker kill.

⁠System Requirements

You will need at least 16GB of RAM to load the model and 4 CPU cores. You will also need at least 10GB of storage for the Docker image and dataset that is downloaded.

I would recommend using the 32GB memory-optimized droplet from DigitalOcean as it provides plenty of RAM and an adequate number of CPU cores to train on. It should cost at most $5, so long as you remember to shut down the droplet once you are done using it.

⁠Retrieving the Models

The generated models can be found in the volume directory under models/output_tflite_graph.tflite. Throughout the training process, checkpoints are generated which can be found under the ckpt directory. Each checkpoint is suffixed with the epoch number it was generated at.

⁠Examples

• Retrain MobileNet v1

docker run --privileged -p 6006:6006 -v $PWD:/tensorflow/models/research/powercell-detection akrantz/powercell-detection:latest

• Retrain MobileNet v2

docker run --privileged -p 6006:6006 $PWD:/tensorflow/models/research/powercell-detection -e NETWORK_TYPE=mobilenet_v2_ssd akrantz/powercell-detection:latest

⁠Options

⁠NETWORK_TYPE

This variable determines which base network to retrain. The valid options are mobilenet_v1_ssd and mobilenet_v2_ssd.

⁠TRAIN_WHOLE_MODEL

This variable determines whether all of the model layers will be retrained. It is false by default, in which only the last few layers are trained. If set to true, it will take roughly 16 hours as opposed to roughly 5 hours. At the moment of writing this, I have not retrained all the layers myself so I cannot say if you will gain better results.

⁠NUM_TRAINING_STEPS

This variable determines the number of training steps (or epochs) to evaluate. It is set to 500 by default.

⁠NUM_EVAL_STEPS

This variable determines the number of evaluation steps to run. It is set to 100 by default.