| 232 | ) |
| 233 | |
| 234 | def update(self, minibatches, unlabeled=None): |
| 235 | device = "cuda" if minibatches[0][0].is_cuda else "cpu" |
| 236 | self.update_count += 1 |
| 237 | all_x = torch.cat([x for x, y in minibatches]) |
| 238 | all_y = torch.cat([y for x, y in minibatches]) |
| 239 | all_z = self.featurizer(all_x) |
| 240 | if self.conditional: |
| 241 | disc_input = all_z + self.class_embeddings(all_y) |
| 242 | else: |
| 243 | disc_input = all_z |
| 244 | disc_out = self.discriminator(disc_input) |
| 245 | disc_labels = torch.cat( |
| 246 | [ |
| 247 | torch.full((x.shape[0],), i, dtype=torch.int64, device=device) |
| 248 | for i, (x, y) in enumerate(minibatches) |
| 249 | ] |
| 250 | ) |
| 251 | |
| 252 | if self.class_balance: |
| 253 | y_counts = F.one_hot(all_y).sum(dim=0) |
| 254 | weights = 1. / (y_counts[all_y] * y_counts.shape[0]).float() |
| 255 | disc_loss = F.cross_entropy(disc_out, disc_labels, reduction='none') |
| 256 | disc_loss = (weights * disc_loss).sum() |
| 257 | else: |
| 258 | disc_loss = F.cross_entropy(disc_out, disc_labels) |
| 259 | |
| 260 | disc_softmax = F.softmax(disc_out, dim=1) |
| 261 | input_grad = autograd.grad( |
| 262 | disc_softmax[:, disc_labels].sum(), [disc_input], create_graph=True |
| 263 | )[0] |
| 264 | grad_penalty = (input_grad**2).sum(dim=1).mean(dim=0) |
| 265 | disc_loss += self.hparams['grad_penalty'] * grad_penalty |
| 266 | |
| 267 | d_steps_per_g = self.hparams['d_steps_per_g_step'] |
| 268 | if (self.update_count.item() % (1 + d_steps_per_g) < d_steps_per_g): |
| 269 | |
| 270 | self.disc_opt.zero_grad() |
| 271 | disc_loss.backward() |
| 272 | self.disc_opt.step() |
| 273 | return {'disc_loss': disc_loss.item()} |
| 274 | else: |
| 275 | all_preds = self.classifier(all_z) |
| 276 | classifier_loss = F.cross_entropy(all_preds, all_y) |
| 277 | gen_loss = (classifier_loss + (self.hparams['lambda'] * -disc_loss)) |
| 278 | self.disc_opt.zero_grad() |
| 279 | self.gen_opt.zero_grad() |
| 280 | gen_loss.backward() |
| 281 | self.gen_opt.step() |
| 282 | return {'gen_loss': gen_loss.item()} |
| 283 | |
| 284 | def predict(self, x): |
| 285 | return self.classifier(self.featurizer(x)) |