diff options
18 files changed, 229 insertions, 153 deletions
@@ -51,8 +51,8 @@ Makefile.config # 1. reference, and not casually committed # 2. custom, and live on their own unless they're deliberated contributed data/* -*model -*_iter_* +models/* +*.caffemodel *.solverstate *.binaryproto *leveldb diff --git a/docs/getting_pretrained_models.md b/docs/getting_pretrained_models.md index 5df2bd4d..70d8f7bd 100644 --- a/docs/getting_pretrained_models.md +++ b/docs/getting_pretrained_models.md @@ -5,30 +5,12 @@ layout: default # Pre-trained models [BVLC](http://bvlc.eecs.berkeley.edu) aims to provide a variety of high quality pre-trained models. -Note that unlike Caffe itself, these models are licensed for **academic research / non-commercial use only**. -If you have any questions, please get in touch with us. +Note that unlike Caffe itself, these models usually have licenses **academic research / non-commercial use only**. -*UPDATE* July 2014: we are actively working on a service for hosting user-uploaded model definition and trained weight files. -Soon, the community will be able to easily contribute different architectures! +## TODO -### ImageNet +Write something about the model zoo. -**Caffe Reference ImageNet Model**: Our reference implementation of an ImageNet model trained on ILSVRC-2012 can be downloaded (232.6MB) by running `examples/imagenet/get_caffe_reference_imagenet_model.sh` from the Caffe root directory. - -- The bundled model is the iteration 310,000 snapshot. -- The best validation performance during training was iteration 313,000 with - validation accuracy 57.412% and loss 1.82328. -- This model obtains a top-1 accuracy 57.4% and a top-5 accuracy 80.4% on the validation set, using just the center crop. (Using the average of 10 crops, (4 + 1 center) * 2 mirror, should obtain a bit higher accuracy) - -**AlexNet**: Our training of the Krizhevsky architecture, which differs from the paper's methodology by (1) not training with the relighting data-augmentation and (2) initializing non-zero biases to 0.1 instead of 1. (2) was found necessary for training, as initialization to 1 gave flat loss. Download the model (243.9MB) by running `examples/imagenet/get_caffe_alexnet_model.sh` from the Caffe root directory. - -- The bundled model is the iteration 360,000 snapshot. -- The best validation performance during training was iteration 358,000 with - validation accuracy 57.258% and loss 1.83948. -- This model obtains a top-1 accuracy 57.1% and a top-5 accuracy 80.2% on the validation set, using just the center crop. (Using the average of 10 crops, (4 + 1 center) * 2 mirror, should obtain a bit higher accuracy) - -**R-CNN (ILSVRC13)**: The pure Caffe instantiation of the [R-CNN](https://github.com/rbgirshick/rcnn) model for ILSVRC13 detection. Download the model (230.8MB) by running `examples/imagenet/get_caffe_rcnn_imagenet_model.sh` from the Caffe root directory. This model was made by transplanting the R-CNN SVM classifiers into a `fc-rcnn` classification layer, provided here as an off-the-shelf Caffe detector. Try the [detection example](http://nbviewer.ipython.org/github/BVLC/caffe/blob/master/examples/detection.ipynb) to see it in action. For the full details, refer to the R-CNN site. *N.B. For research purposes, make use of the official R-CNN package and not this example.* - -### Auxiliary Data +## Auxiliary Data Additionally, you will probably eventually need some auxiliary data (mean image, synset list, etc.): run `data/ilsvrc12/get_ilsvrc_aux.sh` from the root directory to obtain it. diff --git a/docs/model_zoo.md b/docs/model_zoo.md new file mode 100644 index 00000000..b4cbeb1d --- /dev/null +++ b/docs/model_zoo.md @@ -0,0 +1,31 @@ +# Caffe Model Zoo + +A caffe model is distributed as a directory containing: +- solver/model prototxt(s) +- model binary file, with .caffemodel extension +- readme.md, containing: + - YAML header: + - model file URL or (torrent magnet link) and MD5 hash + - Caffe commit hash use to train this model + - [optional] github gist id + - license type or text + - main body: free-form description/details +- helpful scripts + +It is up to the user where to host the model file. +Dropbox or their own server are both fine. + +We provide scripts: + +- publish_model_as_gist.sh: uploads non-binary files in the model directory as a Github Gist and returns the id. If gist id is already part of the readme, then updates existing gist. +- download_model_from_gist.sh <gist_id>: downloads the non-binary files from a Gist. +- download_model_binary.py: downloads the .caffemodel from the URL specified in readme. + +The Gist is a good format for distribution because it can contain multiple files, is versionable, and has in-browser syntax highlighting and markdown rendering. + +The existing models distributed with Caffe can stay bundled with Caffe, so I am re-working them all into this format. +All relevant examples will be updated to start with `cd models/model_of_interest && ../scripts/download_model_binary.sh`. + +## Tasks + +- get the imagenet example to work with the new prototxt location diff --git a/examples/imagenet/get_caffe_alexnet_model.sh b/examples/imagenet/get_caffe_alexnet_model.sh deleted file mode 100755 index 7312ed93..00000000 --- a/examples/imagenet/get_caffe_alexnet_model.sh +++ /dev/null @@ -1,28 +0,0 @@ -#!/usr/bin/env sh -# This scripts downloads the caffe reference imagenet model -# for ilsvrc image classification and deep feature extraction - -MODEL=caffe_alexnet_model -CHECKSUM=29eb495b11613825c1900382f5286963 - -if [ -f $MODEL ]; then - echo "Model already exists. Checking md5..." - os=`uname -s` - if [ "$os" = "Linux" ]; then - checksum=`md5sum $MODEL | awk '{ print $1 }'` - elif [ "$os" = "Darwin" ]; then - checksum=`cat $MODEL | md5` - fi - if [ "$checksum" = "$CHECKSUM" ]; then - echo "Model checksum is correct. No need to download." - exit 0 - else - echo "Model checksum is incorrect. Need to download again." - fi -fi - -echo "Downloading..." - -wget http://dl.caffe.berkeleyvision.org/$MODEL examples/imagenet/$MODEL - -echo "Done. Please run this command again to verify that checksum = $CHECKSUM." diff --git a/examples/imagenet/get_caffe_rcnn_imagenet_model.sh b/examples/imagenet/get_caffe_rcnn_imagenet_model.sh deleted file mode 100755 index 9a8d0a15..00000000 --- a/examples/imagenet/get_caffe_rcnn_imagenet_model.sh +++ /dev/null @@ -1,28 +0,0 @@ -#!/usr/bin/env sh -# This scripts downloads the Caffe R-CNN ImageNet -# for ILSVRC13 detection. - -MODEL=caffe_rcnn_imagenet_model -CHECKSUM=42c1556d2d47a9128c4a90e0a9c5341c - -if [ -f $MODEL ]; then - echo "Model already exists. Checking md5..." - os=`uname -s` - if [ "$os" = "Linux" ]; then - checksum=`md5sum $MODEL | awk '{ print $1 }'` - elif [ "$os" = "Darwin" ]; then - checksum=`cat $MODEL | md5` - fi - if [ "$checksum" = "$CHECKSUM" ]; then - echo "Model checksum is correct. No need to download." - exit 0 - else - echo "Model checksum is incorrect. Need to download again." - fi -fi - -echo "Downloading..." - -wget http://dl.caffe.berkeleyvision.org/$MODEL examples/imagenet/$MODEL - -echo "Done. Please run this command again to verify that checksum = $CHECKSUM." diff --git a/examples/imagenet/get_caffe_reference_imagenet_model.sh b/examples/imagenet/get_caffe_reference_imagenet_model.sh deleted file mode 100755 index f687ebfa..00000000 --- a/examples/imagenet/get_caffe_reference_imagenet_model.sh +++ /dev/null @@ -1,28 +0,0 @@ -#!/usr/bin/env sh -# This scripts downloads the caffe reference imagenet model -# for ilsvrc image classification and deep feature extraction - -MODEL=caffe_reference_imagenet_model -CHECKSUM=af678f0bd3cdd2437e35679d88665170 - -if [ -f $MODEL ]; then - echo "Model already exists. Checking md5..." - os=`uname -s` - if [ "$os" = "Linux" ]; then - checksum=`md5sum $MODEL | awk '{ print $1 }'` - elif [ "$os" = "Darwin" ]; then - checksum=`cat $MODEL | md5` - fi - if [ "$checksum" = "$CHECKSUM" ]; then - echo "Model checksum is correct. No need to download." - exit 0 - else - echo "Model checksum is incorrect. Need to download again." - fi -fi - -echo "Downloading..." - -wget http://dl.caffe.berkeleyvision.org/$MODEL examples/imagenet/$MODEL - -echo "Done. Please run this command again to verify that checksum = $CHECKSUM." diff --git a/examples/imagenet/readme.md b/examples/imagenet/readme.md index b4a3110e..6b18407d 100644 --- a/examples/imagenet/readme.md +++ b/examples/imagenet/readme.md @@ -6,18 +6,16 @@ include_in_docs: true priority: 1 --- -Yangqing's Recipe on Brewing ImageNet -===================================== +Brewing ImageNet +================ - "All your braincells are belong to us." - - Caffeine - -We are going to describe a reference implementation for the approach first proposed by Krizhevsky, Sutskever, and Hinton in their [NIPS 2012 paper](http://books.nips.cc/papers/files/nips25/NIPS2012_0534.pdf). Since training the whole model takes some time and energy, we provide a model, trained in the same way as we describe here, to help fight global warming. If you would like to simply use the pretrained model, check out the [Pretrained ImageNet](../../getting_pretrained_models.html) page. *Note that the pretrained model is for academic research / non-commercial use only*. +We are going to describe a reference implementation for the approach first proposed by Krizhevsky, Sutskever, and Hinton in their [NIPS 2012 paper](http://books.nips.cc/papers/files/nips25/NIPS2012_0534.pdf). +Since training the whole model takes some time and energy, we provide a model, trained in the same way as we describe here, to help fight global warming. +If you would like to simply use the pretrained model, check out the [Pretrained ImageNet](../../getting_pretrained_models.html) page. +*Note that the pretrained model is for academic research / non-commercial use only*. To clarify, by ImageNet we actually mean the ILSVRC12 challenge, but you can easily train on the whole of ImageNet as well, just with more disk space, and a little longer training time. -(If you don't get the quote, visit [Yann LeCun's fun page](http://yann.lecun.com/ex/fun/). - Data Preparation ---------------- @@ -100,11 +98,13 @@ We all experience times when the power goes out, or we feel like rewarding ourse ./resume_training.sh -where in the script `caffe_imagenet_train_1000.solverstate` is the solver state snapshot that stores all necessary information to recover the exact solver state (including the parameters, momentum history, etc). +where in the script `imagenet_train_1000.solverstate` is the solver state snapshot that stores all necessary information to recover the exact solver state (including the parameters, momentum history, etc). Parting Words ------------- -Hope you liked this recipe! Many researchers have gone further since the ILSVRC 2012 challenge, changing the network architecture and/or finetuning the various parameters in the network. The recent ILSVRC 2013 challenge suggests that there are quite some room for improvement. **Caffe allows one to explore different network choices more easily, by simply writing different prototxt files** - isn't that exciting? +Hope you liked this recipe! +Many researchers have gone further since the ILSVRC 2012 challenge, changing the network architecture and/or finetuning the various parameters in the network. +**Caffe allows one to explore different network choices more easily, by simply writing different prototxt files** - isn't that exciting? -And since now you have a trained network, check out how to use it: [Running Pretrained ImageNet](../../getting_pretrained_models.html). This time we will use Python, but if you have wrappers for other languages, please kindly send a pull request! +And since now you have a trained network, check out how to use it with the Python interface: [Running Pretrained ImageNet](../../getting_pretrained_models.html). diff --git a/examples/imagenet/alexnet_deploy.prototxt b/models/bvlc_alexnet/alexnet_deploy.prototxt index d010753f..d010753f 100644 --- a/examples/imagenet/alexnet_deploy.prototxt +++ b/models/bvlc_alexnet/alexnet_deploy.prototxt diff --git a/examples/imagenet/alexnet_solver.prototxt b/models/bvlc_alexnet/alexnet_solver.prototxt index 94bda7f3..0d61fed4 100644 --- a/examples/imagenet/alexnet_solver.prototxt +++ b/models/bvlc_alexnet/alexnet_solver.prototxt @@ -1,4 +1,4 @@ -net: "examples/imagenet/alexnet_train_val.prototxt" +net: "models/bvlc_alexnet/alexnet_train_val.prototxt" test_iter: 1000 test_interval: 1000 base_lr: 0.01 @@ -10,5 +10,5 @@ max_iter: 450000 momentum: 0.9 weight_decay: 0.0005 snapshot: 10000 -snapshot_prefix: "examples/imagenet/caffe_alexnet" +snapshot_prefix: "models/bvlc_alexnet/caffe_alexnet_train" solver_mode: GPU diff --git a/examples/imagenet/alexnet_train_val.prototxt b/models/bvlc_alexnet/alexnet_train_val.prototxt index 3fa46773..69b8916d 100644 --- a/examples/imagenet/alexnet_train_val.prototxt +++ b/models/bvlc_alexnet/alexnet_train_val.prototxt @@ -34,6 +34,8 @@ layers { layers { name: "conv1" type: CONVOLUTION + bottom: "data" + top: "conv1" blobs_lr: 1 blobs_lr: 2 weight_decay: 1 @@ -51,8 +53,6 @@ layers { value: 0 } } - bottom: "data" - top: "conv1" } layers { name: "relu1" @@ -63,28 +63,30 @@ layers { layers { name: "norm1" type: LRN + bottom: "conv1" + top: "norm1" lrn_param { local_size: 5 alpha: 0.0001 beta: 0.75 } - bottom: "conv1" - top: "norm1" } layers { name: "pool1" type: POOLING + bottom: "norm1" + top: "pool1" pooling_param { pool: MAX kernel_size: 3 stride: 2 } - bottom: "norm1" - top: "pool1" } layers { name: "conv2" type: CONVOLUTION + bottom: "pool1" + top: "conv2" blobs_lr: 1 blobs_lr: 2 weight_decay: 1 @@ -103,8 +105,6 @@ layers { value: 0.1 } } - bottom: "pool1" - top: "conv2" } layers { name: "relu2" @@ -115,28 +115,30 @@ layers { layers { name: "norm2" type: LRN + bottom: "conv2" + top: "norm2" lrn_param { local_size: 5 alpha: 0.0001 beta: 0.75 } - bottom: "conv2" - top: "norm2" } layers { name: "pool2" type: POOLING + bottom: "norm2" + top: "pool2" pooling_param { pool: MAX kernel_size: 3 stride: 2 } - bottom: "norm2" - top: "pool2" } layers { name: "conv3" type: CONVOLUTION + bottom: "pool2" + top: "conv3" blobs_lr: 1 blobs_lr: 2 weight_decay: 1 @@ -154,8 +156,6 @@ layers { value: 0 } } - bottom: "pool2" - top: "conv3" } layers { name: "relu3" @@ -166,6 +166,8 @@ layers { layers { name: "conv4" type: CONVOLUTION + bottom: "conv3" + top: "conv4" blobs_lr: 1 blobs_lr: 2 weight_decay: 1 @@ -184,8 +186,6 @@ layers { value: 0.1 } } - bottom: "conv3" - top: "conv4" } layers { name: "relu4" @@ -196,6 +196,8 @@ layers { layers { name: "conv5" type: CONVOLUTION + bottom: "conv4" + top: "conv5" blobs_lr: 1 blobs_lr: 2 weight_decay: 1 @@ -214,8 +216,6 @@ layers { value: 0.1 } } - bottom: "conv4" - top: "conv5" } layers { name: "relu5" @@ -226,17 +226,19 @@ layers { layers { name: "pool5" type: POOLING + bottom: "conv5" + top: "pool5" pooling_param { pool: MAX kernel_size: 3 stride: 2 } - bottom: "conv5" - top: "pool5" } layers { name: "fc6" type: INNER_PRODUCT + bottom: "pool5" + top: "fc6" blobs_lr: 1 blobs_lr: 2 weight_decay: 1 @@ -252,8 +254,6 @@ layers { value: 0.1 } } - bottom: "pool5" - top: "fc6" } layers { name: "relu6" @@ -264,15 +264,17 @@ layers { layers { name: "drop6" type: DROPOUT + bottom: "fc6" + top: "fc6" dropout_param { dropout_ratio: 0.5 } - bottom: "fc6" - top: "fc6" } layers { name: "fc7" type: INNER_PRODUCT + bottom: "fc6" + top: "fc7" blobs_lr: 1 blobs_lr: 2 weight_decay: 1 @@ -288,8 +290,6 @@ layers { value: 0.1 } } - bottom: "fc6" - top: "fc7" } layers { name: "relu7" @@ -300,15 +300,17 @@ layers { layers { name: "drop7" type: DROPOUT + bottom: "fc7" + top: "fc7" dropout_param { dropout_ratio: 0.5 } - bottom: "fc7" - top: "fc7" } layers { name: "fc8" type: INNER_PRODUCT + bottom: "fc7" + top: "fc8" blobs_lr: 1 blobs_lr: 2 weight_decay: 1 @@ -324,8 +326,6 @@ layers { value: 0 } } - bottom: "fc7" - top: "fc8" } layers { name: "accuracy" diff --git a/models/bvlc_alexnet/readme.md b/models/bvlc_alexnet/readme.md new file mode 100644 index 00000000..20c393ff --- /dev/null +++ b/models/bvlc_alexnet/readme.md @@ -0,0 +1,25 @@ +--- +name: BVLC AlexNet Model +caffemodel: bvlc_alexnet.caffemodel +caffemodel_url: http://dl.caffe.berkeleyvision.org/bvlc_alexnet.caffemodel +license: non-commercial +sha1: 9116a64c0fbe4459d18f4bb6b56d647b63920377 +caffe_commit: 709dc15af4a06bebda027c1eb2b3f3e3375d5077 +--- + +This model is a replication of the model described in the [AlexNet](http://papers.nips.cc/paper/4824-imagenet-classification-with-deep-convolutional-neural-networks) publication. + +Differences: +- not training with the relighting data-augmentation; +- initializing non-zero biases to 0.1 instead of 1 (found necessary for training, as initialization to 1 gave flat loss). + +The bundled model is the iteration 360,000 snapshot. +The best validation performance during training was iteration 358,000 with validation accuracy 57.258% and loss 1.83948. +This model obtains a top-1 accuracy 57.1% and a top-5 accuracy 80.2% on the validation set, using just the center crop. +(Using the average of 10 crops, (4 + 1 center) * 2 mirror, should obtain a bit higher accuracy.) + +## License + +The data used to train this model comes from the ImageNet project, which distributes its database to researchers who agree to a following term of access: +"Researcher shall use the Database only for non-commercial research and educational purposes." +Accordingly, this model is distributed under a non-commercial license. diff --git a/examples/imagenet/imagenet_deploy.prototxt b/models/bvlc_reference_caffenet/caffenet_deploy.prototxt index 4e494f42..4e494f42 100644 --- a/examples/imagenet/imagenet_deploy.prototxt +++ b/models/bvlc_reference_caffenet/caffenet_deploy.prototxt diff --git a/examples/imagenet/imagenet_solver.prototxt b/models/bvlc_reference_caffenet/caffenet_solver.prototxt index 5b5be4bb..d52e11a5 100644 --- a/examples/imagenet/imagenet_solver.prototxt +++ b/models/bvlc_reference_caffenet/caffenet_solver.prototxt @@ -1,4 +1,4 @@ -net: "examples/imagenet/imagenet_train_val.prototxt" +net: "models/bvlc_reference_caffenet/caffenet_train_val.prototxt" test_iter: 1000 test_interval: 1000 base_lr: 0.01 @@ -10,5 +10,5 @@ max_iter: 450000 momentum: 0.9 weight_decay: 0.0005 snapshot: 10000 -snapshot_prefix: "examples/imagenet/caffe_imagenet" +snapshot_prefix: "models/bvlc_reference_caffenet/caffenet_train" solver_mode: GPU diff --git a/examples/imagenet/imagenet_train_val.prototxt b/models/bvlc_reference_caffenet/caffenet_train_val.prototxt index d6d64073..d6d64073 100644 --- a/examples/imagenet/imagenet_train_val.prototxt +++ b/models/bvlc_reference_caffenet/caffenet_train_val.prototxt diff --git a/models/bvlc_reference_caffenet/readme.md b/models/bvlc_reference_caffenet/readme.md new file mode 100644 index 00000000..1fbdbe12 --- /dev/null +++ b/models/bvlc_reference_caffenet/readme.md @@ -0,0 +1,26 @@ +--- +name: BVLC CaffeNet Model +caffemodel: bvlc_reference_caffenet.caffemodel +caffemodel_url: http://dl.caffe.berkeleyvision.org/bvlc_reference_caffenet.caffemodel +license: non-commercial +sha1: 4c8d77deb20ea792f84eb5e6d0a11ca0a8660a46 +caffe_commit: 709dc15af4a06bebda027c1eb2b3f3e3375d5077 +--- + +This model is the result of following the Caffe [instructions](http://caffe.berkeleyvision.org/gathered/examples/imagenet.html) on training an ImageNet model. +This model is a replication of the model described in the [AlexNet](http://papers.nips.cc/paper/4824-imagenet-classification-with-deep-convolutional-neural-networks) publication with some differences: + +- not training with the relighting data-augmentation; +- the order of pooling and normalization layers is switched (in CaffeNet, pooling is done before normalization). + + +This model is snapshot of iteration 310,000. +The best validation performance during training was iteration 313,000 with validation accuracy 57.412% and loss 1.82328. +This model obtains a top-1 accuracy 57.4% and a top-5 accuracy 80.4% on the validation set, using just the center crop. +(Using the average of 10 crops, (4 + 1 center) * 2 mirror, should obtain a bit higher accuracy still.) + +## License + +The data used to train this model comes from the ImageNet project, which distributes its database to researchers who agree to a following term of access: +"Researcher shall use the Database only for non-commercial research and educational purposes." +Accordingly, this model is distributed under a non-commercial license. diff --git a/examples/imagenet/rcnn_imagenet_deploy.prototxt b/models/bvlc_reference_rcnn_ilsvrc13/rcnn_imagenet_deploy.prototxt index ef75a0a5..ef75a0a5 100644 --- a/examples/imagenet/rcnn_imagenet_deploy.prototxt +++ b/models/bvlc_reference_rcnn_ilsvrc13/rcnn_imagenet_deploy.prototxt diff --git a/models/bvlc_reference_rcnn_ilsvrc13/readme.md b/models/bvlc_reference_rcnn_ilsvrc13/readme.md new file mode 100644 index 00000000..fb8f26d1 --- /dev/null +++ b/models/bvlc_reference_rcnn_ilsvrc13/readme.md @@ -0,0 +1,20 @@ +--- +name: BVLC Reference RCNN ILSVRC13 Model +caffemodel: bvlc_reference_rcnn_ilsvrc13.caffemodel +caffemodel_url: http://dl.caffe.berkeleyvision.org/bvlc_reference_rcnn_ilsvrc13.caffemodel +license: non-commercial +sha1: bdd8abb885819cba5e2fe1eb36235f2319477e64 +caffe_commit: a7e397abbda52c0b90323c23ab95bdeabee90a98 +--- + +The pure Caffe instantiation of the [R-CNN](https://github.com/rbgirshick/rcnn) model for ILSVRC13 detection. +This model was made by transplanting the R-CNN SVM classifiers into a `fc-rcnn` classification layer, provided here as an off-the-shelf Caffe detector. +Try the [detection example](http://nbviewer.ipython.org/github/BVLC/caffe/blob/master/examples/detection.ipynb) to see it in action. + +*N.B. For research purposes, make use of the official R-CNN package and not this example.* + +## License + +The data used to train this model comes from the ImageNet project, which distributes its database to researchers who agree to a following term of access: +"Researcher shall use the Database only for non-commercial research and educational purposes." +Accordingly, this model is distributed under a non-commercial license. diff --git a/scripts/download_model_binary.py b/scripts/download_model_binary.py new file mode 100755 index 00000000..48e9015f --- /dev/null +++ b/scripts/download_model_binary.py @@ -0,0 +1,76 @@ +#!/usr/bin/env python +import os +import sys +import time +import yaml +import urllib +import hashlib +import argparse + +required_keys = ['caffemodel', 'caffemodel_url', 'sha1'] + + +def reporthook(count, block_size, total_size): + """ + From http://blog.moleculea.com/2012/10/04/urlretrieve-progres-indicator/ + """ + global start_time + if count == 0: + start_time = time.time() + return + duration = time.time() - start_time + progress_size = int(count * block_size) + speed = int(progress_size / (1024 * duration)) + percent = int(count * block_size * 100 / total_size) + sys.stdout.write("\r...%d%%, %d MB, %d KB/s, %d seconds passed" % + (percent, progress_size / (1024 * 1024), speed, duration)) + sys.stdout.flush() + + +def parse_readme_frontmatter(dirname): + readme_filename = os.path.join(dirname, 'readme.md') + with open(readme_filename) as f: + lines = [line.strip() for line in f.readlines()] + top = lines.index('---') + bottom = lines[top + 1:].index('---') + frontmatter = yaml.load('\n'.join(lines[top + 1:bottom])) + assert all(key in frontmatter for key in required_keys) + return dirname, frontmatter + + +def valid_dirname(dirname): + try: + return parse_readme_frontmatter(dirname) + except Exception as e: + print('ERROR: {}'.format(e)) + raise argparse.ArgumentTypeError( + 'Must be valid Caffe model directory with a correct readme.md') + + +if __name__ == '__main__': + parser = argparse.ArgumentParser( + description='Download trained model binary.') + parser.add_argument('dirname', type=valid_dirname) + args = parser.parse_args() + + # A tiny hack: the dirname validator also returns readme YAML frontmatter. + dirname = args.dirname[0] + frontmatter = args.dirname[1] + model_filename = os.path.join(dirname, frontmatter['caffemodel']) + + # Closure-d function for checking SHA1. + def model_checks_out(filename=model_filename, sha1=frontmatter['sha1']): + with open(filename, 'r') as f: + return hashlib.sha1(f.read()).hexdigest() == sha1 + + # Check if model exists. + if os.path.exists(model_filename) and model_checks_out(): + print("Model already exists.") + sys.exit(0) + + # Download and verify model. + urllib.urlretrieve( + frontmatter['caffemodel_url'], model_filename, reporthook) + if not model_checks_out(): + print('ERROR: model did not download correctly! Run this again.') + sys.exit(1) |