Allow copy TRAIN.WEIGHTS to checkpoint folder

Summary:
Add TRAIN.COPY_WEIGHTS (default False):

When setting True, training will copy TRAIN.WEIGHTS to checkpoint folder and treat it as a candidate checkpoint. This is useful if we want to resume a checkpoint from another job (instead of using ImageNet pre-training).

Reviewed By: rbgirshick

Differential Revision: D8710711

fbshipit-source-id: 9c011227867c52eaaeb40e3c2679cff1cfccdc66

detectron/core/config.py

@@ -189,6 +189,9 @@ __C.TRAIN.FREEZE_CONV_BODY = False
 # output directory
 __C.TRAIN.AUTO_RESUME = True
 
+# Training will copy TRAIN.WEIGHTS and treat it as a candidate checkpoint
+__C.TRAIN.COPY_WEIGHTS = False
+
 # Add StopGrad at a specified stage so the bottom layers are frozen
 __C.TRAIN.FREEZE_AT = 2
 

detectron/utils/train.py

@@ -28,6 +28,7 @@ from __future__ import division
 from __future__ import print_function
 from __future__ import unicode_literals
 
+from shutil import copyfile
 import cv2  # NOQA (Must import before importing caffe2 due to bug in cv2)
 import logging
 import numpy as np
@@ -109,6 +110,12 @@ def create_model():
             logger.info('model_final.pkl exists; no need to train!')
             return None, None, None, {'final': final_path}, output_dir
 
+        if cfg.TRAIN.COPY_WEIGHTS:
+            copyfile(
+                weights_file,
+                os.path.join(output_dir, os.path.basename(weights_file)))
+            logger.info('Copy {} to {}'.format(weights_file, output_dir))
+
         # Find the most recent checkpoint (highest iteration number)
         files = os.listdir(output_dir)
         for f in files: