Merge pull request #118 from lincc-frameworks/train_script

Train script

configs/solo/solo_swin_DC2_new.py → configs/solo/solo_swin.py

@@ -4,6 +4,7 @@
 
 from omegaconf import OmegaConf
 import numpy as np
+import os
 # ---------------------------------------------------------------------------- #
 # Local variables and metadata
 # ---------------------------------------------------------------------------- #
@@ -67,19 +68,46 @@
     box_predictor.test_score_thresh = 0.5
     box_predictor.test_nms_thresh = 0.3
 
-#The ImageNet1k pretrained weights file
+#The ImageNet1k pretrained weights file.  Update to your own path
 train.init_checkpoint = "/home/shared/hsc/detectron2/projects/ViTDet/model_final_246a82.pkl"
 
 optimizer.lr = 0.001
 dataloader.test.mapper = loaders.DictMapper
 dataloader.train.mapper = loaders.DictMapper
+dataloader.epoch=epoch
 
 # ---------------------------------------------------------------------------- #
 #Change for different data sets
 reader = DC2ImageReader()
 dataloader.imagereader = reader
+
+# Key_mapper will take a metadatadict and return the key that the imagereader will use to read in the corresponding image
+# Implemented so that if you move images on the disk or save as a different format, you don't have to change filepaths in the metadata
+# Mostly, one can just have it return the filename key in the dictionary
+def key_mapper(dataset_dict):
+    '''
+    args
+        dataset_dict: [dict]
+            A dictionary of metadata
+    
+    returns
+        fn: str
+            The filepath to the corresponding image
+    
+    '''
+    filename = dataset_dict["filename"]
+    base = os.path.basename(filename)
+    dirpath = "../tests/deepdisc/test_data/dc2/"
+    fn = os.path.join(dirpath, base)
+    return fn
+
+
+dataloader.key_mapper = key_mapper
+
 # ---------------------------------------------------------------------------- #
-dataloader.epoch=epoch
+
+
+
 
 
 # ---------------------------------------------------------------------------- #

configs/solo/solo_swin_hsc.py

@@ -0,0 +1,132 @@
+""" This is a demo "solo config" file for use in solo_test_run_transformers.py.
+
+This uses template configs cascade_mask_rcnn_swin_b_in21k_50ep and yaml_style_defaults."""
+
+from omegaconf import OmegaConf
+import numpy as np
+import os
+# ---------------------------------------------------------------------------- #
+# Local variables and metadata
+# ---------------------------------------------------------------------------- #
+epoch=2
+bs=1
+metadata = OmegaConf.create() 
+metadata.classes = ["star", "galaxy"]
+
+numclasses = len(metadata.classes)
+
+# ---------------------------------------------------------------------------- #
+# Standard config (this has always been the LazyConfig/.py-style config)
+# ---------------------------------------------------------------------------- #
+# Get values from templates
+from ..COCO.cascade_mask_rcnn_swin_b_in21k_50ep import dataloader, model, train, lr_multiplier, optimizer
+import deepdisc.model.loaders as loaders
+from deepdisc.data_format.augment_image import train_augs
+from deepdisc.data_format.image_readers import HSCImageReader
+
+# Overrides
+dataloader.augs = train_augs
+dataloader.train.total_batch_size = bs
+
+model.proposal_generator.anchor_generator.sizes = [[8], [16], [32], [64], [128]]
+model.roi_heads.num_classes = numclasses
+model.roi_heads.batch_size_per_image = 512
+
+model.roi_heads.num_classes = numclasses
+model.roi_heads.batch_size_per_image = 512
+
+
+# ---------------------------------------------------------------------------- #
+#Change for different data sets
+
+#This is the number of color channels in the images
+model.backbone.bottom_up.in_chans = 3         
+
+# ---------------------------------------------------------------------------- #
+model.proposal_generator.nms_thresh = 0.3
+
+for box_predictor in model.roi_heads.box_predictors:
+    box_predictor.test_topk_per_image = 2000
+    box_predictor.test_score_thresh = 0.5
+    box_predictor.test_nms_thresh = 0.3
+
+#The ImageNet1k pretrained weights file.  Update to your own path
+train.init_checkpoint = "/home/shared/hsc/detectron2/projects/ViTDet/model_final_246a82.pkl"
+#train.init_checkpoint = "/home/shared/hsc/AAS/Swin_astrolupton_new.pth"
+
+optimizer.lr = 0.001
+dataloader.test.mapper = loaders.DictMapper
+dataloader.train.mapper = loaders.DictMapper
+dataloader.epoch=epoch
+
+# ---------------------------------------------------------------------------- #
+#Change for different data sets
+reader = HSCImageReader(norm='lupton')
+dataloader.imagereader = reader
+
+# Key_mapper will take a metadatadict and return the key that the imagereader will use to read in the corresponding image
+# Implemented so that if you move images on the disk or save as a different format, you don't have to change filepaths in the metadata
+# Mostly, one can just have it return the filename key in the dictionary
+def key_mapper(dataset_dict):
+    '''
+    args
+        dataset_dict: [dict]
+            A dictionary of metadata
+    
+    returns
+        fn: str
+            The filepath to the corresponding image
+    
+    '''
+    filenames = []
+    for b in ['G','R','I']:
+        fn = dataset_dict[f"filename_{b}"]
+        base = os.path.basename(fn)
+        dirpath = "../tests/deepdisc/test_data/hsc/"
+        fn = os.path.join(dirpath, base)   
+        filenames.append(fn)
+    return filenames
+
+
+dataloader.key_mapper = key_mapper
+
+# ---------------------------------------------------------------------------- #
+
+
+
+
+
+# ---------------------------------------------------------------------------- #
+# Yaml-style config (was formerly saved as a .yaml file, loaded to cfg_loader)
+# ---------------------------------------------------------------------------- #
+# Get values from template
+from .yacs_style_defaults import MISC, DATALOADER, DATASETS, GLOBAL, INPUT, MODEL, SOLVER, TEST
+
+# Overrides
+SOLVER.IMS_PER_BATCH = bs
+
+DATASETS.TRAIN = "astro_train"
+DATASETS.TEST = "astro_val"
+
+SOLVER.BASE_LR = 0.001
+SOLVER.CLIP_GRADIENTS.ENABLED = True
+# Type of gradient clipping, currently 2 values are supported:
+# - "value": the absolute values of elements of each gradients are clipped
+# - "norm": the norm of the gradient for each parameter is clipped thus
+#   affecting all elements in the parameter
+SOLVER.CLIP_GRADIENTS.CLIP_TYPE = "norm"
+# Maximum absolute value used for clipping gradients
+# Floating point number p for L-p norm to be used with the "norm"
+# gradient clipping type; for L-inf, please specify .inf
+SOLVER.CLIP_GRADIENTS.NORM_TYPE = 5.0
+
+
+e1 = epoch * 15
+e2 = epoch * 25
+e3 = epoch * 30
+efinal = epoch * 50
+
+SOLVER.STEPS = [e1,e2,e3]  # do not decay learning rate for retraining
+SOLVER.LR_SCHEDULER_NAME = "WarmupMultiStepLR"
+SOLVER.WARMUP_ITERS = 0
+SOLVER.MAX_ITER = efinal  # for DefaultTrainer

scripts/README.md

@@ -0,0 +1,23 @@
+
+## Training script:  
+
+This directory contains the script used to run the full training, ```run_model.py```  This will work for .py configs, but not .yacs configs (yet)
+
+Run the script with ```python run_model.py --cfgfile $path_to_config --train-metadata $path_to_train_dicts --eval-metadata $path_to_eval_dicts --num-gpus $ngpu --run-name $name_of_run --output-dir $path_to_output.```  
+
+You can test this with the double/single_test.json files in ```/tests/deepdisc/test_data/dc2/``` and the config in ```/configs/solo/solo_swin.py```  You should download the pre-trained weights [here](https://dl.fbaipublicfiles.com/detectron2/ViTDet/COCO/cascade_mask_rcnn_swin_b_in21k/f342979038/model_final_246a82.pkl)
+
+Other pre-trained models using transformers available [here](https://github.com/facebookresearch/detectron2/tree/main/projects/ViTDet)
+
+The command line options are explained below  
+
+- cfgfile: The configuration file used to build the model, learning rate optimizer, trainer, and dataloaders.
+- train-metadata: The training data as a list of dicts stored in json format.  The dicts should have the "instance detection/segmentation" keys specified in the [detectron2 repo](https://detectron2.readthedocs.io/en/latest/tutorials/datasets.html)
+- eval-metadata: The same as the training metadata, but for the evaluation set.
+- num-gpus: The number of gpus used to train the model.  Must be a multiple of the batch size specified in the config
+- run-name: A string prefix that will be used to save the outputs of the script such as model weights and loss curves
+- output-dir: The directory to save the outputs  
+
+After training, inference can be done by loading a predictor (as in the demo notebook) with ```predictor = return_predictor_transformer(cfg)```.  You can use the same config that was used in training, but change the train.init_checkpoint path to the newly saved model.
+
+