(
self,
model_arch: str,
input_image_shape: tuple = (3, 224, 224),
agent_feature_dim: int = None,
global_feature_dim: int = None,
roi_feature_size: tuple = (7, 7),
roi_layer_key: str = "layer4",
output_activation=nn.ReLU,
use_rotated_roi=True
)
| 670 | """Use RoI Align to crop map feature for each agent""" |
| 671 | |
| 672 | def __init__( |
| 673 | self, |
| 674 | model_arch: str, |
| 675 | input_image_shape: tuple = (3, 224, 224), |
| 676 | agent_feature_dim: int = None, |
| 677 | global_feature_dim: int = None, |
| 678 | roi_feature_size: tuple = (7, 7), |
| 679 | roi_layer_key: str = "layer4", |
| 680 | output_activation=nn.ReLU, |
| 681 | use_rotated_roi=True |
| 682 | ) -> None: |
| 683 | super(RasterizeROIEncoder, self).__init__() |
| 684 | encoder = RasterizedMapEncoder( |
| 685 | model_arch=model_arch, |
| 686 | input_image_shape=input_image_shape, |
| 687 | feature_dim=global_feature_dim, |
| 688 | ) |
| 689 | feat_nodes = { |
| 690 | 'map_model.layer1': 'layer1', |
| 691 | 'map_model.layer2': 'layer2', |
| 692 | 'map_model.layer3': 'layer3', |
| 693 | 'map_model.layer4': 'layer4', |
| 694 | 'map_model.fc': "final" |
| 695 | } |
| 696 | self.encoder_heads = create_feature_extractor(encoder, feat_nodes) |
| 697 | |
| 698 | self.roi_layer_key = roi_layer_key |
| 699 | roi_scale = encoder.feature_scales()[roi_layer_key] |
| 700 | roi_channel = encoder.feature_channels()[roi_layer_key] |
| 701 | if use_rotated_roi: |
| 702 | self.roi_align = RotatedROIAlign( |
| 703 | roi_feature_size, roi_scale=roi_scale) |
| 704 | else: |
| 705 | self.roi_align = RoIAlign( |
| 706 | output_size=roi_feature_size, |
| 707 | spatial_scale=roi_scale, |
| 708 | sampling_ratio=-1, |
| 709 | aligned=True |
| 710 | ) |
| 711 | |
| 712 | self.activation = output_activation() |
| 713 | if agent_feature_dim is not None: |
| 714 | self.agent_net = nn.Sequential( |
| 715 | nn.AdaptiveAvgPool2d((1, 1)), # [B, C, 1, 1] |
| 716 | nn.Flatten(start_dim=1), |
| 717 | nn.Linear(roi_channel, agent_feature_dim), |
| 718 | self.activation |
| 719 | ) |
| 720 | else: |
| 721 | self.agent_net = nn.Identity() |
| 722 | |
| 723 | def forward(self, map_inputs: torch.Tensor, rois: torch.Tensor): |
| 724 | """ |
nothing calls this directly
no test coverage detected