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()
