def model(num_keypoints: int,
num_classes: int = 2,
backbone: Optional[nn.Module] = None,
remove_internal_transforms: bool = True,
pretrained: bool = True,
**keypoint_rcnn_kwargs) -> nn.Module:
"""KeypointRCNN model implemented by torchvision.
# Arguments
num_keypoints: Number of keypoints (e.g. 17 in case of COCO).
num_classes: Number of classes (including background).
backbone: Backbone model to use. Defaults to a resnet50_fpn model.
remove_internal_transforms: The torchvision model internally applies transforms
like resizing and normalization, but we already do this at the `Dataset` level,
so it's safe to remove those internal transforms.
pretrained: Argument passed to `keypointrcnn_resnet50_fpn` if `backbone is None`.
By default it is set to True: this is generally used when training a new model (transfer learning).
`pretrained = False` is used during inference (prediction) for cases where the users have their own pretrained weights.
**keypoint_rcnn_kwargs: Keyword arguments that internally are going to be passed to
`torchvision.models.detection.keypoint_rcnn.KeypointRCNN`.
# Returns
A Pytorch `nn.Module`.
"""
if backbone is None:
model = keypointrcnn_resnet50_fpn(pretrained=pretrained,
pretrained_backbone=pretrained,
**keypoint_rcnn_kwargs)
in_channels = model.roi_heads.keypoint_predictor.kps_score_lowres.in_channels
model.roi_heads.keypoint_predictor = KeypointRCNNPredictor(
in_channels, num_keypoints)
if num_classes != 2:
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(
in_features, num_classes)
resnet_fpn.patch_param_groups(model.backbone)
else:
model = KeypointRCNN(backbone,
num_classes=num_classes,
num_keypoints=num_keypoints,
**keypoint_rcnn_kwargs)
patch_rcnn_param_groups(model=model)
if remove_internal_transforms:
remove_internal_model_transforms(model)
return model
def test_keypoint_rcnn(self):
model = keypoint_rcnn.keypointrcnn_resnet50_fpn(
pretrained=False,
pretrained_backbone=False,
min_size=200,
max_size=300)
images, test_images = _get_test_images()
self.run_test(model, (images, ), rtol=1e-3, atol=1e-5)
self.run_test(
model,
(images, ),
input_names=["images_tensors"],
output_names=["outputs1", "outputs2", "outputs3", "outputs4"],
dynamic_axes={"images_tensors": [0, 1, 2]},
rtol=1e-3,
atol=1e-5,
)
dummy_images = [torch.ones(3, 100, 100) * 0.3]
self.run_test(
model,
(images, ),
additional_test_inputs=[(images, ), (test_images, ),
(dummy_images, )],
input_names=["images_tensors"],
output_names=["outputs1", "outputs2", "outputs3", "outputs4"],
dynamic_axes={"images_tensors": [0, 1, 2]},
rtol=5e-3,
atol=1e-5,
)
self.run_test(
model,
(dummy_images, ),
additional_test_inputs=[(dummy_images, ), (test_images, )],
input_names=["images_tensors"],
output_names=["outputs1", "outputs2", "outputs3", "outputs4"],
dynamic_axes={"images_tensors": [0, 1, 2]},
rtol=5e-3,
atol=1e-5,
)
# prepare data loaders
train_loader, valid_loader = prepare_train_valid_loaders(
trainset, valid_size, batch_size)
# print(train_loader)
# for batch in train_loader:
# print(batch)
# plt.imshow(batch['image'][0].permute(2, 1, 0))
# plt.title(batch['answer'][0])
# plt.show()
# break
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
pretrained_model = keypoint_rcnn.keypointrcnn_resnet50_fpn(pretrained=True)
pretrained_model.eval()
model = MLP(input_size=batch_size,
output_size=10,
hidden_layers=[10, 10],
pretrained_model_flag=False,
drop_p=0.1)
model = model.to(device)
model = model.double()
criterion = torch.nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.003)
train(train_loader,
valid_loader,
model,
criterion,
optimizer,
def get_model():
return keypointrcnn_resnet50_fpn(pretrained=True,
progress=True,
num_classes=2,
pretrained_backbone=True)
