def train_model(trainset_name: str, learning_rate: int, num_iteration: int, batch_per_image: int, num_classes: int): cfg = get_cfg() cfg.MODEL.DEVICE = "cpu" cfg.merge_from_file( model_zoo.get_config_file( "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")) cfg.DATASETS.TRAIN = (trainset_name, ) cfg.DATASETS.TEST = () cfg.DATALOADER.NUM_WORKERS = 2 cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url( "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml" ) # Let training initialize from model zoo cfg.SOLVER.IMS_PER_BATCH = 2 cfg.SOLVER.BASE_LR = learning_rate # pick a good LR cfg.SOLVER.MAX_ITER = num_iteration # 300 iterations seems good enough for this toy dataset; you will need to train longer for a practical dataset cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = batch_per_image # faster, and good enough for this toy dataset (default: 512) cfg.MODEL.ROI_HEADS.NUM_CLASSES = num_classes # (see https://detectron2.readthedocs.io/tutorials/datasets.html#update-the-config-for-new-datasets) os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=False) trainer.train() # Inference should use the config with parameters that are used in training # cfg now already contains everything we've set previously. We changed it a little bit for inference: cfg.MODEL.WEIGHTS = os.path.join( cfg.OUTPUT_DIR, "model_final.pth") # path to the model we just trained cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 # set a custom testing threshold predictor = DefaultPredictor(cfg) return predictor
def train_detectron2(): dataset_storage = { 'synthetic': generate_synthetic_datasets(), 'real': generate_real_datasets() } cfg = get_cfg() cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")) cfg.DATASETS.TRAIN = (SYNTHETIC_DATASET_NAME[MODES[0]], REAL_DATASET_NAME[MODES[0]]) cfg.DATASETS.TEST = () cfg.DATALOADER.NUM_WORKERS = 2 cfg.MODEL.WEIGHTS = "detectron2://COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl" cfg.SOLVER.IMS_PER_BATCH = 2 cfg.SOLVER.BASE_LR = 0.00025 # pick a good LR cfg.SOLVER.MAX_ITER = 500 cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 # faster, and good enough for this toy dataset (default: 512) cfg.MODEL.ROI_HEADS.NUM_CLASSES = len(dataset_storage["synthetic"][MODES[0]]['unit_classes']) cfg.OUTPUT_DIR = str(OUTPUT_PATH) os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=False) trainer.train() return cfg
def train(train_dir, name_data, json_dir, config, resume_status, iteration, batch, lr): cfg = get_cfg() cfg.merge_from_file("./configs/" + config) cfg.SOLVER.IMS_PER_BATCH = 4 cfg.SOLVER.BASE_LR = lr cfg.SOLVER.WARMUP_ITERS = 1200 cfg.SOLVER.MAX_ITER = iteration cfg.SOLVER.GAMMA = 0.05 cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = batch cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 cfg.DATASETS.TRAIN = (name_data, ) cfg.DATASETS.TEST = () try: register_coco_instances(name_data, {}, json_dir, train_dir) # Train data except ValueError: print("Data already registerd. Continue.") if resume_status: cfg.MODEL.WEIGHTS = os.path.join('output', 'model_final.pth') os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=resume_status) trainer.train()
def task_b_MOTS_training(model_name, model_file): #model_name = model_name + '_inference' print('Running task B for model', model_name) SAVE_PATH = os.path.join('./results_week_5_task_b', model_name) os.makedirs(SAVE_PATH, exist_ok=True) # Load model and configuration print('Loading Model') cfg = get_cfg() cfg.merge_from_file(model_zoo.get_config_file(model_file)) cfg.DATASETS.TRAIN = ('MOTS_train',) cfg.DATASETS.TEST = ('KITTIMOTS_val',) cfg.DATALOADER.NUM_WORKERS = 0 cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 cfg.OUTPUT_DIR = SAVE_PATH cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(model_file) cfg.SOLVER.IMS_PER_BATCH = 4 cfg.SOLVER.BASE_LR = 0.00025 cfg.SOLVER.MAX_ITER = 200 cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 256 cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3 cfg.TEST.SCORE_THRESH = 0.5 # Training print('Training') trainer = DefaultTrainer(cfg) val_loss = ValidationLoss(cfg) trainer.register_hooks([val_loss]) trainer._hooks = trainer._hooks[:-2] + trainer._hooks[-2:][::-1] trainer.resume_or_load(resume=False) trainer.train() # Evaluation print('Evaluating') evaluator = COCOEvaluator('KITTIMOTS_val', cfg, False, output_dir=SAVE_PATH) trainer.model.load_state_dict(val_loss.weights) trainer.test(cfg, trainer.model, evaluators=[evaluator]) print('Plotting losses') draw_loss(cfg, cfg.SOLVER.MAX_ITER, model_name, SAVE_PATH) # Qualitative results: visualize some results print('Getting qualitative results') predictor = DefaultPredictor(cfg) predictor.model.load_state_dict(trainer.model.state_dict()) inputs = kitti_val() #inputs = inputs[:20] + inputs[-20:] inputs = inputs[220:233] + inputs[1995:2100] for i, input in enumerate(inputs): file_name = input['file_name'] print('Prediction on image ' + file_name) img = cv2.imread(file_name) outputs = predictor(img) v = Visualizer( img[:, :, ::-1], metadata=MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=0.8, instance_mode=ColorMode.IMAGE) v = v.draw_instance_predictions(outputs['instances'].to('cpu')) cv2.imwrite(os.path.join(SAVE_PATH, 'Inference_' + model_name + '_inf_' + str(i) + '.png'), v.get_image()[:, :, ::-1])
def train(train_flag,resume_load=False): # trainer= Trainer(cfg) trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume_load) if train_flag: trainer.train() return trainer
def Train(): register_coco_instances( "custom", {}, "/home/lsc/datasets/butterfly/Annotations/train.json", "/home/lsc/datasets/butterfly/TrainData/JPEGImages") custom_metadata = MetadataCatalog.get("custom") dataset_dicts = DatasetCatalog.get("custom") cfg = get_cfg() cfg.merge_from_file( "/home/lsc/detectron2/configs/COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml" ) cfg.DATASETS.TRAIN = ("custom", ) cfg.DATASETS.TEST = () cfg.DATALOADER.NUM_WORKERS = 8 cfg.MODEL.WEIGHTS = 'detectron2://COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x/139173657/model_final_68b088.pkl' cfg.SOLVER.IMS_PER_BATCH = 2 cfg.SOLVER.BASE_LR = 0.02 cfg.SOLVER.MAX_ITER = (500) cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = (128) cfg.MODEL.ROI_HEADS.NUM_CLASSES = 94 os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=False) trainer.train()
def train_iobjectspy_voc(train_data_path, train_config_path, weight_path, max_iter, out_dir, register_train_name, ml_set_tracking_path, experiment_id, ml_experiment_tag): cfg = get_cfg() cfg.merge_from_file(train_config_path) cfg.DATASETS.TRAIN = (register_train_name, ) cfg.DATASETS.TEST = () # no metrics implemented for this dataset cfg.MODEL.WEIGHTS = weight_path # initialize from model zoo if max_iter == -1: pass else: cfg.SOLVER.MAX_ITER = max_iter num_class = get_class_num(train_data_path) cfg.MODEL.ROI_HEADS.NUM_CLASSES = num_class # get classes from sda cfg.OUTPUT_DIR = out_dir os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=True) trainer.train() try: import mlflow as ml # 设置mlflow ml.set_tracking_uri(ml_set_tracking_path) # 通过设置不同的实验id来管理实验,建议这一层级为项目名称,比如:iobjectspy_faster_rcnn_dota ml.set_experiment(experiment_id) # 通过设置 ml.set_tag('experiment_id', ml_experiment_tag) ml.log_param('lr', cfg.SOLVER.BASE_LR) ml.log_param('max_iter', cfg.SOLVER.MAX_ITER) ml.log_param('epoch', cfg.SOLVER.IMS_PER_BATCH) except: pass
def train(*args): prepare_dataset() # D2 configuration cfg = get_cfg() cfg.merge_from_file( model_zoo.get_config_file( "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")) cfg.DATASETS.TRAIN = ("balloon_train", ) cfg.DATASETS.TEST = () cfg.DATALOADER.NUM_WORKERS = 2 cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url( "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml" ) # Let training initialize from model zoo cfg.SOLVER.IMS_PER_BATCH = args[ 0] # number ims_per_batch should be divisible by number of workers. D2 assertion. cfg.SOLVER.BASE_LR = 0.00025 # pick a good LR cfg.SOLVER.MAX_ITER = 300 # 300 iterations seems good enough for this toy dataset; you may need to train longer for a practical dataset cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 # faster, and good enough for this toy dataset (default: 512) cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 # only has one class (ballon) cfg.OUTPUT_DIR = os.environ[ 'SM_OUTPUT_DATA_DIR'] # TODO check that this config works fine trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=False) trainer.train()
def Train(): register_coco_instances( "custom", {}, "datasets/coco/annotations/instances_train2017.json", "datasets/coco/train2017") custom_metadata = MetadataCatalog.get("custom") dataset_dicts = DatasetCatalog.get("custom") for d in random.sample(dataset_dicts, 3): img = cv2.imread(d["file_name"]) visualizer = Visualizer(img[:, :, ::-1], metadata=custom_metadata, scale=1) vis = visualizer.draw_dataset_dict(d) cv2.imshow('Sample', vis.get_image()[:, :, ::-1]) cv2.waitKey() cfg = get_cfg() cfg.merge_from_file( "configs/COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml") cfg.DATASETS.TRAIN = ("custom", ) cfg.DATASETS.TEST = () cfg.DATALOADER.NUM_WORKERS = 4 cfg.MODEL.WEIGHTS = 'model_final_3c3198.pkl' cfg.SOLVER.IMS_PER_BATCH = 2 cfg.SOLVER.BASE_LR = 0.0001 cfg.SOLVER.MAX_ITER = (150000) cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = (128) cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=False) trainer.train()
def main(): args = parse_args() with open(args.config, "r") as f: config = yaml.safe_load(f) if "classes" not in config: raise Exception("Could not find class names") n_classes = len(config["classes"]) classes = config["classes"] for d in ["train"]: DatasetCatalog.register("custom_" + d, lambda d=d: get_annotated_dataset(args.annotator_root, args.data_folders)) MetadataCatalog.get("custom_" + d).set(thing_classes=classes) custom_metadata = MetadataCatalog.get("custom_train") cfg = get_cfg() cfg.merge_from_file(args.model_config) cfg.DATASETS.TRAIN = ("custom_train",) cfg.DATASETS.TEST = () # no metrics implemented for this dataset cfg.DATALOADER.NUM_WORKERS = 2 cfg.MODEL.WEIGHTS = args.initial_weights cfg.SOLVER.IMS_PER_BATCH = 2 cfg.SOLVER.BASE_LR = 0.00025 cfg.SOLVER.MAX_ITER = args.max_iter cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 # faster, and good enough for this toy dataset cfg.MODEL.ROI_HEADS.NUM_CLASSES = len(classes) os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=False) trainer.train()
def model_configuration(model_url, learning_rate, max_iter): model_a = os.path.join("COCO-InstanceSegmentation", model_url) cfg = get_cfg() cfg.merge_from_file(model_zoo.get_config_file(model_a)) cfg.DATASETS.TRAIN = ("kitti-mots", ) cfg.DATASETS.TEST = () cfg.DATALOADER.NUM_WORKERS = 2 cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url( "COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml" ) # Let training initialize from model zoo cfg.SOLVER.IMS_PER_BATCH = 2 cfg.SOLVER.BASE_LR = learning_rate cfg.SOLVER.MAX_ITER = max_iter cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512 cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2 # cfg.INPUT.MASK_FORMAT = 'rle' # cfg.INPUT.MASK_FORMAT='bitmask' os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=False) print('Training...') trainer.train() # EVALUATION print('Evaluating...') evaluator = COCOEvaluator("kitti-mots", cfg, False, output_dir="./Search/") val_loader = build_detection_test_loader(cfg, "kitti-mots") results = inference_on_dataset(trainer.model, val_loader, evaluator) print(results) return results
def task_b(model_name, model_file): save_path = Path("output/task_b") / model_name os.makedirs(save_path, exist_ok=True) cfg = base_cfg(model_file, save_path) cfg.DATALOADER.NUM_WORKERS = 0 cfg.SOLVER.IMS_PER_BATCH = 4 cfg.SOLVER.BASE_LR = 0.00025 cfg.SOLVER.MAX_ITER = 1000 cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 256 cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3 cfg.TEST.SCORE_THRESH = 0.5 trainer = DefaultTrainer(cfg) val_loss = ValidationLoss(cfg) trainer.register_hooks([val_loss]) trainer._hooks = trainer._hooks[:-2] + trainer._hooks[-2:][::-1] trainer.resume_or_load(resume=False) trainer.train() evaluator = COCOEvaluator("kitti-mots-val", cfg, False, output_dir=save_path) trainer.model.load_state_dict(val_loss.weights) trainer.test(cfg, trainer.model, evaluators=[evaluator]) plot_validation_loss(cfg, cfg.SOLVER.MAX_ITER, model_name, save_path) get_qualitative_results(cfg, save_path)
def main(args): dataset_train = "faces_train" dataset_validation = "faces_validation" cfg = setup(args) cfg.DATASETS.TRAIN = (dataset_train, ) ### TODO: ### - Test with validation set. ### # cfg.DATASETS.TEST = (dataset_validation,) cfg.DATASETS.TEST = () train_data = args.train_data train_annots = args.train_annotations validation_data = args.validation_data validation_annots = args.validation_annotations DatasetCatalog.register(dataset_train, lambda: get_dataset_dict(train_annots, train_data)) MetadataCatalog.get(dataset_train).set(thing_classes=["face"]) DatasetCatalog.register( dataset_validation, lambda: get_dataset_dict(validation_annots, validation_data)) MetadataCatalog.get(dataset_validation).set(thing_classes=["face"]) faces_metadata = MetadataCatalog.get(dataset_train) os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=False) trainer.train()
def train(args, mode, _appcfg): name = "hmd" for subset in ["train", "val"]: metadata = load_and_register_dataset(name, subset, _appcfg) cfg = config.get_cfg() cfg.merge_from_file( "/aimldl-cod/external/detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml" ) cfg.DATASETS.TRAIN = ("hmd_train", "hmd_val") cfg.DATASETS.TEST = () cfg.DATALOADER.NUM_WORKERS = 2 cfg.MODEL.WEIGHTS = "detectron2://COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl" # initialize from model zoo # cfg.MODEL.WEIGHTS = "/codehub/apps/detectron2/release/model_final.pth" cfg.SOLVER.IMS_PER_BATCH = 2 cfg.SOLVER.BASE_LR = 0.00025 # cfg.SOLVER.MAX_ITER = 350000 # 300 iterations seems good enough, but you can certainly train longer # cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 # faster, and good enough for this toy dataset cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512 # faster, and good enough for this toy dataset cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3 # only has one class (ballon) os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=False) trainer.train()
def Train(): os.environ["CUDA_VISIBLE_DEVICES"] = "0" # register train dataset register_coco_instances("custom", {}, "datasets/traindata/midv500_coco.json", "datasets/traindata/") custom_metadata = MetadataCatalog.get("custom") dataset_dicts = DatasetCatalog.get("custom") # set cfg cfg = get_cfg() cfg.merge_from_file( "configs/COCO-InstanceSegmentation/mask_rcnn_R_50_DC5_1x.yaml") cfg.DATASETS.TRAIN = ("custom", ) cfg.DATASETS.TEST = () cfg.DATALOADER.NUM_WORKERS = 4 cfg.MODEL.WEIGHTS = 'model_final_maskrcnn_dc5.pkl' cfg.SOLVER.IMS_PER_BATCH = 2 cfg.SOLVER.BASE_LR = 0.02 cfg.SOLVER.MAX_ITER = (1000) cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = (512) cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 # training os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=False) trainer.train()
def train(args): dataset_name = "dataset_train" register_coco_instances(dataset_name, {}, args.annotations_path, args.images_dir) cfg = get_cfg() if args.type.lower() == "maskrcnn": cfg.merge_from_file( model_zoo.get_config_file( "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")) else: cfg.merge_from_file( model_zoo.get_config_file( "COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml")) setKeypoints(dataset_name) cfg.MODEL.ROI_KEYPOINT_HEAD.NUM_KEYPOINTS = 14 cfg.DATASETS.TRAIN = (dataset_name, ) cfg.DATASETS.TEST = () cfg.INPUT.MASK_FORMAT = 'bitmask' cfg.DATALOADER.NUM_WORKERS = 2 setWeights(args, cfg) cfg.SOLVER.IMS_PER_BATCH = 2 cfg.SOLVER.BASE_LR = 0.00025 # pick a good LR cfg.SOLVER.MAX_ITER = 500 # 300 iterations seems good enough for this toy dataset; you may need to train longer for a practical dataset cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 # faster, and good enough for this toy dataset (default: 512) setNumClasses(cfg) os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=False) trainer.train()
def main(args): # first regiest dataset I will use register_coco_instances("my_dataset_train", {}, "training.json", "../data/ade20k/full_data/images/training/") register_coco_instances("my_dataset_val", {}, "validation.json", "../data/ade20k/full_data/images/validation/") # this is just a default cfg files cfg = get_cfg() # accordinig to different yaml file, it will change cfg files accordiningly cfg.merge_from_file("../detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml") # This is some task specific changes I made for training ade20k dataset cfg.DATASETS.TRAIN = ("my_dataset_train",) cfg.DATASETS.TEST = () # no metrics implemented for this dataset cfg.DATALOADER.NUM_WORKERS = 2 cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml") cfg.MODEL.ROI_HEADS.NUM_CLASSES = 150 # 150 classes cfg.SOLVER.IMS_PER_BATCH = 16 # this is default one os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) # I highly suggest read source code of DefaultTrainer again, if you forget why you did this. trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=False) trainer.train()
def prepare_for_training(N_iter, output_dir, train_dataset_name, N_classes, start_training=False, gpu_avail=True, model_type="COCO-Detection/faster_rcnn_R_50_C4_1x.yaml"): cfg = get_cfg() cfg.merge_from_file(model_zoo.get_config_file(model_type)) cfg.OUTPUT_DIR = output_dir cfg.DATASETS.TRAIN = (train_dataset_name,) cfg.DATASETS.TEST = () cfg.DATALOADER.NUM_WORKERS = 2 cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(model_type) # Let training initialize from model zoo cfg.SOLVER.IMS_PER_BATCH = 2 cfg.SOLVER.BASE_LR = 0.00025 # pick a good LR cfg.SOLVER.GAMMA = 0.99 #lr decay cfg.SOLVER.STEPS = list(range(1000, N_iter, 1000)) #(decay steps,) cfg.SOLVER.WARMUP_ITERS = 500 #warmup steps cfg.SOLVER.MAX_ITER = N_iter # 300 iterations seems good enough for this toy dataset; you may need to train longer for a practical dataset cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 # faster, and good enough for this toy dataset (default: 512) cfg.MODEL.ROI_HEADS.NUM_CLASSES = N_classes # 4 classes if not gpu_avail: cfg.MODEL.DEVICE = 'cpu' os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=False) if start_training: trainer.train() return trainer, cfg
def train_coco_data(self, coco_json): dataset_name = "mask_train_data" DatasetCatalog.register( dataset_name, lambda: load_coco_json(json_file=coco_json, image_root=self.train_data_path)) MetadataCatalog.get(dataset_name).set( json_file=coco_json, image_root=self.train_data_path, evaluator_type="coco", thing_classes=["rightmask"], thing_dataset_id_to_contiguous_id={1: 0}) cfg = get_cfg() cfg.merge_from_file( "/home/detectron2/configs/COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml" ) cfg.DATASETS.TRAIN = (dataset_name, ) cfg.DATASETS.TEST = (dataset_name, ) cfg.DATALOADER.NUM_WORKERS = 2 cfg.MODEL.WEIGHTS = "/home/detectron2/train_data/model_final_280758.pkl" cfg.SOLVER.IMS_PER_BATCH = 2 cfg.SOLVER.BASE_LR = 0.01 # 学习率 cfg.SOLVER.MAX_ITER = 300 # 最大迭代次数 cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) print("模型存储路径" + cfg.OUTPUT_DIR) trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=False) trainer.train() # 开始训练
def train_model(): cfg = get_cfg() cfg.merge_from_file(model_zoo.get_config_file( "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")) cfg.DATASETS.TRAIN = ("balloon_train",) cfg.DATASETS.TEST = () cfg.DATALOADER.NUM_WORKERS = 2 # Let training initialize from model zoo cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url( "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml") cfg.SOLVER.IMS_PER_BATCH = 2 cfg.SOLVER.BASE_LR = 0.00025 # pick a good LR # 300 iterations seems good enough for this toy dataset; # you may need to train longer for a practical dataset cfg.SOLVER.MAX_ITER = 300 # faster, and good enough for this toy dataset (default: 512) cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 # only has one class (ballon) cfg.OUTPUT_DIR = gcfg.get_ou_dir os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=False) trainer.train()
def train(config, data_path): """Train the Mask-RCNN for the given configuration and the given data""" register_data(data_path, prefix='yeast_cells_') os.makedirs(config.OUTPUT_DIR, exist_ok=True) trainer = DefaultTrainer(config) trainer.resume_or_load(resume=True) trainer.train() return trainer
def main(setupoptions: SetupOptions): baseStr = 'VerticalNanowires_noSn_16' maskType = setDatasetAndMetadata(baseStr, setupoptions) configurator = setConfigurator(setupoptions, baseStr, maskType) trainer = DefaultTrainer(configurator) trainer.resume_or_load(resume=setupoptions.continueTraining) trainer.train()
def train_model(dataset): # Export the dataset to COCO format export_file, image_dir = export_dataset(dataset) # Register it as a COCO dataset in the Detectron2 framework try: register_coco_instances('my_dataset', {}, export_file, image_dir) except: print('Dataset was already registered') dataset_dicts = load_coco_json(export_file, image_dir) MetadataCatalog.get('my_dataset').set( thing_classes=[c['name'] for c in dataset.categories]) segments_metadata = MetadataCatalog.get('my_dataset') print(segments_metadata) # Configure the training run cfg = get_cfg() cfg.merge_from_file( model_zoo.get_config_file( 'COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml')) cfg.DATASETS.TRAIN = ('my_dataset', ) cfg.DATASETS.TEST = () cfg.INPUT.MASK_FORMAT = 'bitmask' cfg.DATALOADER.NUM_WORKERS = 2 cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url( 'COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml' ) # Let training initialize from model zoo cfg.SOLVER.IMS_PER_BATCH = 4 # 4 cfg.SOLVER.BASE_LR = 0.00025 # pick a good LR cfg.SOLVER.MAX_ITER = 6000 # 300 iterations seems good enough for this toy dataset; you may need to train longer for a practical dataset cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512 # faster, and good enough for this toy dataset (default: 512) cfg.MODEL.ROI_HEADS.NUM_CLASSES = len( dataset.categories) # number of categories # cfg.MODEL.DEVICE = 'cuda' print('Max iter is ', cfg.SOLVER.MAX_ITER) # Start the training os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=False) trainer.train() # Return the model cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, 'model_final.pth') cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 # set the testing threshold for this model cfg.DATASETS.TEST = ('my_dataset', ) cfg.TEST.DETECTIONS_PER_IMAGE = 1000 built_model = build_model(cfg) # returns a torch.nn.Module DetectionCheckpointer(built_model).load( cfg.MODEL.WEIGHTS) #capture trained model checkpointer = DetectionCheckpointer( built_model, save_dir="/content/gdrive/My Drive/Colab Notebooks") checkpointer.save("model_final") # save to output/model_999.pth predictor = DefaultPredictor(cfg) model = Model(predictor) return model
def main(args): # Register DeepSports dataset register_coco_instances("deepsports_train", {}, os.path.join(args.dataset_path,"train.json"), os.path.join(args.dataset_path,"train")) register_coco_instances("deepsports_val", {}, os.path.join(args.dataset_path,"val.json"), os.path.join(args.dataset_path,"val")) register_coco_instances("deepsports_test", {}, os.path.join(args.dataset_path,"test_nolabels.json"), os.path.join(args.dataset_path,"test")) # Load config file and adjust where needed cfg = get_cfg() cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")) cfg.INPUT.MASK_FORMAT = 'bitmask' cfg.DATASETS.TRAIN = ("deepsports_train",) cfg.DATASETS.TEST = ("deepsports_val",) cfg.DATALOADER.NUM_WORKERS = 4 cfg.SOLVER.IMS_PER_BATCH = args.bs cfg.SOLVER.BASE_LR = args.lr cfg.SOLVER.MAX_ITER = args.iter cfg.SOLVER.STEPS = [] cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2 # human (1), ball (2) cfg.MODEL.WEIGHTS = None # transfer learning is prohibited! print(cfg) # Create output dir os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) if args.predict: # load model and weights model = build_model(cfg) DetectionCheckpointer(model).load(args.predict) # make dataloader test_loader = build_detection_test_loader(cfg, "deepsports_test") # predict outputs = [] model.eval() with torch.no_grad(): for batch in tqdm(test_loader): pred = model(batch) pred = batch_to_dict(pred, batch) # convert pred to serializable format outputs.extend(pred) with open('prediction.json', 'w') as fout: json.dump(outputs, fout) # exit sys.exit() # Train trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=False) trainer.train() evaluator = COCOEvaluator("deepsports_val", ("segm",), False, output_dir="./output/") val_loader = build_detection_test_loader(cfg, "deepsports_val") print(inference_on_dataset(model, val_loader, evaluator))
def train(self, dataset_path, max_iterations=30000): DatasetCatalog.register('training_dataset', lambda : data.to_coco(dataset_path)) self.cfg.DATASETS.TRAIN = ('training_dataset',) self.cfg.DATASETS.TEST = () self.cfg.SOLVER.MAX_ITER = max_iterations trainer = DefaultTrainer(self.cfg) trainer.resume_or_load(resume=False) trainer.train()
def train(self, training_frames, network, get_images): if (training_frames <= 0): raise ValueError( "The number of input frames must be bigger than 0") self.cfg.OUTPUT_DIR = (f'../datasets/detectron2/{network}') self.generate_datasets(training_frames, get_images) retinanet_path = "COCO-Detection/retinanet_R_101_FPN_3x.yaml" faster_rcnn_path = "COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml" if (detectron2.__version__ == "0.1"): if network == 'faster_rcnn': self.cfg.merge_from_file( pkg_resources.resource_filename( "detectron2.model_zoo", os.path.join("configs", faster_rcnn_path))) self.cfg.MODEL.WEIGHTS = model_zoo.ModelZooUrls.get( faster_rcnn_path) if network == 'retinanet': self.cfg.merge_from_file( pkg_resources.resource_filename( "detectron2.model_zoo", os.path.join("configs", retinanet_path))) self.cfg.MODEL.WEIGHTS = model_zoo.ModelZooUrls.get( retinanet_path) else: if network == 'faster_rcnn': self.cfg.merge_from_file( model_zoo.get_config_file(faster_rcnn_path)) self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url( faster_rcnn_path) if network == 'retinanet': self.cfg.merge_from_file( model_zoo.get_config_file(retinanet_path)) self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url( retinanet_path) self.cfg.DATASETS.TRAIN = ('train_set', ) self.cfg.DATASETS.TEST = ('val_set', ) self.cfg.DATALOADER.NUM_WORKERS = 1 self.cfg.SOLVER.IMS_PER_BATCH = 1 self.cfg.SOLVER.BASE_LR = 0.001 self.cfg.SOLVER.MAX_ITER = 2000 self.cfg.SOLVER.STEPS = (1000, 2000) self.cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 self.cfg.MODEL.DEVICE = 'cuda' if not os.path.isfile( os.path.join(self.cfg.OUTPUT_DIR, 'model_final.pth')): os.makedirs(self.cfg.OUTPUT_DIR, exist_ok=True) trainer = DefaultTrainer(self.cfg) trainer.resume_or_load(resume=False) trainer.train()
def main(): # for d in ["train", "val"]: thing_classes = np.load('./dataset/thing_classes.npy', allow_pickle=True) for d in ["trainval"]: DatasetCatalog.register( "aihub_" + d, lambda d=d: get_dataset_dicts(f"dataset/{d}_dataframe.csv")) MetadataCatalog.get("aihub_" + d).set(thing_classes=list(thing_classes)) aihub_metadata = MetadataCatalog.get("aihub_trainval") cfg = get_cfg() cfg.merge_from_file( model_zoo.get_config_file("COCO-Detection/retinanet_R_50_FPN_3x.yaml")) cfg.DATASETS.TRAIN = ("aihub_trainval", ) cfg.DATASETS.TEST = () cfg.DATALOADER.NUM_WORKERS = 12 cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url( "COCO-Detection/retinanet_R_50_FPN_3x.yaml") cfg.SOLVER.IMS_PER_BATCH = 2 cfg.SOLVER.BASE_LR = 0.0001 cfg.SOLVER.MAX_ITER = 3000000 cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 32 cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1031 # RetinaNet cfg.MODEL.RETINANET.NUM_CLASSES = 1031 cfg.MODEL.RETINANET.IN_FEATURES = ["p3", "p4", "p5", "p6", "p7"] cfg.MODEL.RETINANET.NUM_CONVS = 4 cfg.MODEL.RETINANET.IOU_THRESHOLDS = [0.4, 0.5] cfg.MODEL.RETINANET.IOU_LABELS = [0, -1, 1] cfg.MODEL.RETINANET.PRIOR_PROB = 0.01 cfg.MODEL.RETINANET.SCORE_THRESH_TEST = 0.05 cfg.MODEL.RETINANET.TOPK_CANDIDATES_TEST = 1000 cfg.MODEL.RETINANET.NMS_THRESH_TEST = 0.5 # Weights on (dx, dy, dw, dh) for normalizing Retinanet anchor regression targets cfg.MODEL.RETINANET.BBOX_REG_WEIGHTS = (1.0, 1.0, 1.0, 1.0) # Loss parameters cfg.MODEL.RETINANET.FOCAL_LOSS_GAMMA = 2.0 cfg.MODEL.RETINANET.FOCAL_LOSS_ALPHA = 0.25 cfg.MODEL.RETINANET.SMOOTH_L1_LOSS_BETA = 0.1 # Options are: "smooth_l1", "giou" cfg.MODEL.RETINANET.BBOX_REG_LOSS_TYPE = "smooth_l1" cfg.MODEL.RPN.POST_NMS_TOPK_TRAIN = 1500 cfg.DATALOADER.ASPECT_RATIO_GROUPING = False # One of BN, SyncBN, FrozenBN, GN # Only supports GN until unshared norm is implemented cfg.MODEL.RETINANET.NORM = "" os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=False) trainer.train()
def train(self): trainer = DefaultTrainer(self.cfg) trainer.resume_or_load(resume=False) if os.path.exists(self.output): path.remove(self.output) os.makedirs(self.output) else: os.makedirs(self.output) trainer.train() return trainer
def d2_run(): train_metadata, val_metadata = register_dataset() cfg = set_up() trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=False) trainer.train() output_dir = cfg['OUTPUT_DIR'] save_dir = os.path.join(output_dir, 'result') if not os.path.exists(save_dir): os.makedirs(save_dir) cfg.MODEL.WEIGHTS = os.path.join(output_dir, "model_final.pth") cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.2 # set the testing threshold for this model cfg.DATASETS.TEST = (regist_val_name, ) predictor = DefaultPredictor(cfg) test_dir = val_images_dir imgs_list = [ os.path.join(test_dir, file_name) for file_name in os.listdir(test_dir) if file_name.endswith(".jpg") or file_name.endswith(".png") or file_name.endswith(".bmp") or file_name.endswith(".jpeg") ] for d in imgs_list: im = cv2.imread(d) outputs = predictor(im) a = outputs["instances"].pred_classes.data.cpu().numpy().tolist( ) if outputs["instances"].get_fields()["pred_boxes"] else None print(a) v = Visualizer(im[:, :, ::-1], metadata=train_metadata, scale=0.9, instance_mode=ColorMode.IMAGE_BW) v = v.draw_instance_predictions(outputs["instances"].to("cpu")) predict_file = os.path.join( save_dir, os.path.splitext(os.path.basename(d))[0] + "_predict.png") cv2.imwrite(predict_file, v.get_image()[:, :, ::-1]) if os.path.exists(predict_file): print("Done: %s" % predict_file) ### evaluate evaluator = COCOEvaluator(regist_val_name, cfg, False, output_dir="./output/") val_loader = build_detection_test_loader(cfg, regist_val_name) my_eval = inference_on_dataset(trainer.model, val_loader, evaluator) print(my_eval) log = ("%s evaluate: \n" % (model_name), my_eval) print(log, file=open(log_file, "a"))
def train_predictor(self): self.cfg.SOLVER.CHECKPOINT_PERIOD = 2000 self.cfg.SOLVER.IMS_PER_BATCH = 2 self.cfg.SOLVER.BASE_LR /= 2. self.predictor = DefaultPredictor(self.cfg) self.cfg.OUTPUT_DIR = '/scratch/users/zzweng/output/coco/classagnostic' print(self.cfg.OUTPUT_DIR) os.makedirs(self.cfg.OUTPUT_DIR, exist_ok=True) trainer = DefaultTrainer(self.cfg) trainer.resume_or_load(resume=False) trainer.train()