(self, backbone, ca_head, c_head, p_head)
| 135 | |
| 136 | class MultiCropWrapper(nn.Module): |
| 137 | def __init__(self, backbone, ca_head, c_head, p_head): |
| 138 | super(MultiCropWrapper, self).__init__() |
| 139 | # disable layers dedicated to ImageNet labels classification |
| 140 | backbone.fc, backbone.head = nn.Identity(), nn.Identity() |
| 141 | self.backbone = backbone |
| 142 | self.ca_head = ca_head |
| 143 | self.c_head = c_head |
| 144 | self.p_head = p_head |
| 145 | |
| 146 | def forward(self, x, head_only=False, loc=True): |
| 147 | output = self.backbone(x) |