(self, x, x_jig=None, mode=0)
| 151 | head=head) |
| 152 | |
| 153 | def forward(self, x, x_jig=None, mode=0): |
| 154 | # mode -- |
| 155 | # 0: normal encoder, |
| 156 | # 1: momentum encoder, |
| 157 | # 2: testing mode |
| 158 | x1, x2 = torch.split(x, [1, 2], dim=1) |
| 159 | feat1 = self.encoder1(x1) |
| 160 | feat2 = self.encoder2(x2) |
| 161 | |
| 162 | if mode == 0: |
| 163 | x1_jig, x2_jig = torch.split(x_jig, [1, 2], dim=1) |
| 164 | feat1_jig = self.encoder1(x1_jig) |
| 165 | feat2_jig = self.encoder2(x2_jig) |
| 166 | |
| 167 | feat1, feat2 = self.head1(feat1), self.head2(feat2) |
| 168 | feat1_jig = self.head1_jig(feat1_jig) |
| 169 | feat2_jig = self.head2_jig(feat2_jig) |
| 170 | feat = torch.cat((feat1, feat2), dim=1) |
| 171 | feat_jig = torch.cat((feat1_jig, feat2_jig), dim=1) |
| 172 | return feat, feat_jig |
| 173 | elif mode == 1: |
| 174 | feat1, feat2 = self.head1(feat1), self.head2(feat2) |
| 175 | return torch.cat((feat1, feat2), dim=1) |
| 176 | else: |
| 177 | return torch.cat((feat1, feat2), dim=1) |
| 178 | |
| 179 | |
| 180 | NAME_TO_FUNC = { |
nothing calls this directly
no outgoing calls
no test coverage detected