A basic image-based rasterized map encoder
| 557 | |
| 558 | |
| 559 | class RasterizedMapEncoder(nn.Module): |
| 560 | """A basic image-based rasterized map encoder""" |
| 561 | |
| 562 | def __init__( |
| 563 | self, |
| 564 | model_arch: str, |
| 565 | input_image_shape: tuple = (3, 224, 224), |
| 566 | feature_dim: int = None, |
| 567 | use_spatial_softmax=False, |
| 568 | spatial_softmax_kwargs=None, |
| 569 | output_activation=nn.ReLU |
| 570 | ) -> None: |
| 571 | super().__init__() |
| 572 | self.model_arch = model_arch |
| 573 | self.num_input_channels = input_image_shape[0] |
| 574 | self._feature_dim = feature_dim |
| 575 | if output_activation is None: |
| 576 | self._output_activation = nn.Identity() |
| 577 | else: |
| 578 | self._output_activation = output_activation() |
| 579 | |
| 580 | # configure conv backbone |
| 581 | if model_arch == "resnet18": |
| 582 | self.map_model = resnet18() |
| 583 | out_h = int(math.ceil(input_image_shape[1] / 32.)) |
| 584 | out_w = int(math.ceil(input_image_shape[2] / 32.)) |
| 585 | self.conv_out_shape = (512, out_h, out_w) |
| 586 | elif model_arch == "resnet50": |
| 587 | self.map_model = resnet50() |
| 588 | out_h = int(math.ceil(input_image_shape[1] / 32.)) |
| 589 | out_w = int(math.ceil(input_image_shape[2] / 32.)) |
| 590 | self.conv_out_shape = (2048, out_h, out_w) |
| 591 | else: |
| 592 | raise NotImplementedError(f"Model arch {model_arch} unknown") |
| 593 | |
| 594 | # configure spatial reduction pooling layer |
| 595 | if use_spatial_softmax: |
| 596 | pooling = SpatialSoftmax( |
| 597 | input_shape=self.conv_out_shape, **spatial_softmax_kwargs) |
| 598 | self.pool_out_dim = int( |
| 599 | np.prod(pooling.output_shape(self.conv_out_shape))) |
| 600 | else: |
| 601 | pooling = nn.AdaptiveAvgPool2d((1, 1)) |
| 602 | self.pool_out_dim = self.conv_out_shape[0] |
| 603 | |
| 604 | self.map_model.conv1 = nn.Conv2d( |
| 605 | self.num_input_channels, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False |
| 606 | ) |
| 607 | self.map_model.avgpool = pooling |
| 608 | if feature_dim is not None: |
| 609 | self.map_model.fc = nn.Linear( |
| 610 | in_features=self.pool_out_dim, out_features=feature_dim) |
| 611 | else: |
| 612 | self.map_model.fc = nn.Identity() |
| 613 | |
| 614 | def output_shape(self, input_shape=None): |
| 615 | if self._feature_dim is not None: |
| 616 | return [self._feature_dim] |