| 183 | self.support_size = hparams['batch_size'] |
| 184 | |
| 185 | def predict(self, x): |
| 186 | batch_size, c, h, w = x.shape |
| 187 | if batch_size % self.support_size == 0: |
| 188 | meta_batch_size = batch_size // self.support_size |
| 189 | support_size = self.support_size |
| 190 | else: |
| 191 | meta_batch_size, support_size = 1, batch_size |
| 192 | context = self.context_net(x) |
| 193 | context = context.reshape((meta_batch_size, support_size, 1, h, w)) |
| 194 | context = context.mean(dim=1) |
| 195 | context = torch.repeat_interleave(context, repeats=support_size, dim=0) |
| 196 | x = torch.cat([x, context], dim=1) |
| 197 | return self.network(x) |
| 198 | |
| 199 | |
| 200 | class AbstractDANN(Algorithm): |