| 134 | return parser |
| 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) |
| 148 | logits = self.ca_head(output, loc) |
| 149 | return [self.c_head(logits[:,:1]), self.p_head(logits[:,1:])] |
| 150 | |
| 151 | def train_dino(args): |
| 152 | # os.environ['MASTER_PORT'] = '30000' |