| 12 | |
| 13 | |
| 14 | class ImageEncoder(nn.Module): |
| 15 | |
| 16 | def __init__( |
| 17 | self, |
| 18 | trunk: nn.Module, |
| 19 | neck: nn.Module, |
| 20 | scalp: int = 0, |
| 21 | ): |
| 22 | super().__init__() |
| 23 | self.trunk = trunk |
| 24 | self.neck = neck |
| 25 | self.scalp = scalp |
| 26 | assert ( |
| 27 | self.trunk.channel_list == self.neck.backbone_channel_list |
| 28 | ), f"Channel dims of trunk and neck do not match. Trunk: {self.trunk.channel_list}, neck: {self.neck.backbone_channel_list}" |
| 29 | |
| 30 | def forward(self, sample: torch.Tensor): |
| 31 | # Forward through backbone |
| 32 | # features, pos = self.neck(self.trunk(sample)) |
| 33 | |
| 34 | # NOTE: use chunk to reduce memory ------------------------------ |
| 35 | features, pos, chunk_size = [], [], 16 |
| 36 | for base_idx in range(0, sample.size(0), chunk_size): |
| 37 | chunk_features, chunk_pos = self.neck(self.trunk(sample[base_idx:base_idx + chunk_size])) |
| 38 | features.append(chunk_features) |
| 39 | pos.append(chunk_pos) |
| 40 | features = [torch.cat([e[i] for e in features]) for i in range(len(features[0]))] |
| 41 | pos = [torch.cat([e[i] for e in pos]) for i in range(len(pos[0]))] |
| 42 | assert features[0].size(0) == pos[0].size(0) == sample.size(0) |
| 43 | # --------------------------------------------------------------- |
| 44 | |
| 45 | if self.scalp > 0: |
| 46 | # Discard the lowest resolution features |
| 47 | features, pos = features[:-self.scalp], pos[:-self.scalp] |
| 48 | |
| 49 | src = features[-1] |
| 50 | output = { |
| 51 | "vision_features": src, |
| 52 | "vision_pos_enc": pos, |
| 53 | "backbone_fpn": features, |
| 54 | } |
| 55 | return output |
| 56 | |
| 57 | |
| 58 | class FpnNeck(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected