def get_faster_rcnn(): # load an instance segmentation model pre-trained pre-trained on COCO model = fasterrcnn_resnet50_fpn(pretrained=True) # get number of input features for the classifier in_features = model.roi_heads.box_predictor.cls_score.in_features # replace the pre-trained head with a new one model.roi_heads.box_predictor = faster_rcnn.FastRCNNPredictor( in_features, 3) return model
def __init__( self, learning_rate: float = 0.0001, num_classes: int = 91, pretrained: bool = False, pretrained_backbone: bool = True, trainable_backbone_layers: int = 3, replace_head: bool = True, **kwargs, ): """ PyTorch Lightning implementation of `Faster R-CNN: Towards Real-Time Object Detection with Region Proposal Networks <https://arxiv.org/abs/1506.01497>`_. Paper authors: Shaoqing Ren, Kaiming He, Ross Girshick, Jian Sun Model implemented by: - `Teddy Koker <https://github.com/teddykoker>` During training, the model expects both the input tensors, as well as targets (list of dictionary), containing: - boxes (`FloatTensor[N, 4]`): the ground truth boxes in `[x1, y1, x2, y2]` format. - labels (`Int64Tensor[N]`): the class label for each ground truh box CLI command:: # PascalVOC python faster_rcnn.py --gpus 1 --pretrained True Args: learning_rate: the learning rate num_classes: number of detection classes (including background) pretrained: if true, returns a model pre-trained on COCO train2017 pretrained_backbone: if true, returns a model with backbone pre-trained on Imagenet trainable_backbone_layers: number of trainable resnet layers starting from final block """ super().__init__() model = fasterrcnn_resnet50_fpn( # num_classes=num_classes, pretrained=pretrained, pretrained_backbone=pretrained_backbone, trainable_backbone_layers=trainable_backbone_layers, ) if replace_head: in_features = model.roi_heads.box_predictor.cls_score.in_features head = faster_rcnn.FastRCNNPredictor(in_features, num_classes) model.roi_heads.box_predictor = head else: assert num_classes == 91, "replace_head must be true to change num_classes" self.model = model self.learning_rate = learning_rate
def __init__(self, args, num_classes=None): self.num_classes = num_classes if num_classes is not None else args.num_classes self.resnet_name = getattr(args, "resnet_name", "resnet50") nn_module = fasterrcnn_resnet_fpn( name=self.resnet_name, min_size=args.min_size, max_size=args.max_size, ) in_features = nn_module.roi_heads.box_predictor.cls_score.in_features nn_module.roi_heads.box_predictor = faster_rcnn.FastRCNNPredictor( in_features, self.num_classes) super().__init__(nn_module, args)
def main(): device = torch.device( 'cuda') if torch.cuda.is_available() else torch.device('cpu') dataset = BoxyDataset("/media/karsten/data_disk/boxy/boxy_images_raw", get_transforms(), constants.TRAIN_LABEL_FILE) dataset_valid = BoxyDataset( "/media/karsten/data_disk/boxy/boxy_images_raw", get_transforms(), constants.VALID_LABEL_FILE) data_loader = torch.utils.data.DataLoader(dataset, batch_size=4, shuffle=True, num_workers=4, collate_fn=utils.collate_fn) data_loader_valid = torch.utils.data.DataLoader( dataset_valid, batch_size=1, shuffle=False, num_workers=4, collate_fn=utils.collate_fn) model = faster_rcnn.fasterrcnn_resnet50_fpn(pretrained=True) num_classes = 2 # vehicle and background. Do I need background? in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = faster_rcnn.FastRCNNPredictor( in_features, num_classes) model.to(device) params = [p for p in model.parameters() if p.requires_grad] optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005) lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1) num_epochs = 10 for epoch in range(num_epochs): train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10) lr_scheduler.step() evaluate(model, data_loader_valid, device=device) torch.save(model, f"{epoch}_trained.pkl") print("Done training")
def __init__(self, args, num_classes=None): self.num_classes = num_classes if num_classes is not None else args.num_classes nn_module = faster_rcnn.maskrcnn_resnet50_fpn( pretrained=args.pretrained, min_size=args.min_size, max_size=args.max_size) in_features = nn_module.roi_heads.box_predictor.cls_score.in_features nn_module.roi_heads.box_predictor = faster_rcnn.FastRCNNPredictor( in_features, self.num_classes) in_features_mask = nn_module.roi_heads.mask_predictor.conv5_mask.in_channels nh = 256 nn_module.roi_heads.mask_predictor = MaskRCNNPredictor( in_features_mask, nh, self.num_classes) super().__init__(nn_module, args)
def _init_test_roi_heads_faster_rcnn(): out_channels = 256 num_classes = 91 box_fg_iou_thresh = 0.5 box_bg_iou_thresh = 0.5 box_batch_size_per_image = 512 box_positive_fraction = 0.25 bbox_reg_weights = None box_score_thresh = 0.05 box_nms_thresh = 0.5 box_detections_per_img = 100 box_roi_pool = ops.MultiScaleRoIAlign(featmap_names=["0", "1", "2", "3"], output_size=7, sampling_ratio=2) resolution = box_roi_pool.output_size[0] representation_size = 1024 box_head = faster_rcnn.TwoMLPHead(out_channels * resolution**2, representation_size) representation_size = 1024 box_predictor = faster_rcnn.FastRCNNPredictor(representation_size, num_classes) return roi_heads.RoIHeads( box_roi_pool, box_head, box_predictor, box_fg_iou_thresh, box_bg_iou_thresh, box_batch_size_per_image, box_positive_fraction, bbox_reg_weights, box_score_thresh, box_nms_thresh, box_detections_per_img, )
def __init__( self, learning_rate: float = 0.0001, num_classes: int = 91, pretrained: bool = False, pretrained_backbone: bool = True, trainable_backbone_layers: int = 3, replace_head: bool = True, **kwargs, ): """ Args: learning_rate: the learning rate num_classes: number of detection classes (including background) pretrained: if true, returns a model pre-trained on COCO train2017 pretrained_backbone: if true, returns a model with backbone pre-trained on Imagenet trainable_backbone_layers: number of trainable resnet layers starting from final block """ super().__init__() model = fasterrcnn_resnet50_fpn( # num_classes=num_classes, pretrained=pretrained, pretrained_backbone=pretrained_backbone, trainable_backbone_layers=trainable_backbone_layers, ) if replace_head: in_features = model.roi_heads.box_predictor.cls_score.in_features head = faster_rcnn.FastRCNNPredictor(in_features, num_classes) model.roi_heads.box_predictor = head else: assert num_classes == 91, "replace_head must be true to change num_classes" self.model = model self.learning_rate = learning_rate
from argparse import ArgumentParser from DataLoader import WheatDataLoader parser = ArgumentParser() parser = pl.Trainer.add_argparse_args(parent_parser=parser) # parser.add_argument('--batch_size', default=32) # parser.add_argument('--learning_rate', default=1e-3, type=float) parser.add_argument('-dir_input') parser.add_argument('-dir_train') parser.add_argument('-num_workers', default=0, type=int) args = parser.parse_args() model = fasterrcnn_resnet50_fpn(pretrained=True) in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = faster_rcnn.FastRCNNPredictor( in_features, num_classes=2) dl = WheatDataLoader() train_dataloader, val_dataloader = dl.get_data_loaders( dir_input=args.dir_input, dir_train=args.dir_train, debug=True, num_workers=args.num_workers) early_stop_callback = EarlyStopping(monitor='validation_mean_precision', min_delta=0.00, patience=3, verbose=True, mode='max') progress = CustomProgressBar()