download pretrain if missing (tusen-ai#254)

Author: RogerChern

detection_train.py

@@ -129,6 +129,9 @@ def train_net(config):
     elif pModel.from_scratch:
         arg_params, aux_params = dict(), dict()
     else:
+        if not os.path.exists("%s-%04d.params" % (pretrain_prefix, pretrain_epoch)):
+            from utils.download_pretrain import download
+            download(pretrain_prefix, pretrain_epoch)
         arg_params, aux_params = load_checkpoint(pretrain_prefix, pretrain_epoch)
 
     if pModel.process_weight is not None:

utils/download_pretrain.py

@@ -0,0 +1,22 @@
+import os
+import urllib.request
+
+
+def report(block_count, block_size, content_size):
+    if block_count % (content_size // block_size // 5) == 1:
+        print("Downloaded %.1f/100" % (block_size * block_count / content_size * 100))
+
+
+def download(prefix, epoch):
+    base_name = prefix.replace("pretrain_model/", "") + "-%04d.params" % epoch
+    save_name = "%s-%04d.params" % (prefix, epoch)
+    base_url = os.environ.get("SIMPLEDET_BASE_URL", "https://1dv.alarge.space/")
+    full_url = base_url + base_name
+
+    try:
+        print("Downloading %s from %s" % (save_name, full_url))
+        urllib.request.urlretrieve(full_url, save_name, report)
+    except Exception as e:
+        print("Fail to download %s. You can mannually download it from %s and put it to %s" % (base_name, full_url, save_name))
+        os.remove(save_name)
+        raise e