(self, minibatches, unlabeled=None)
| 322 | return result |
| 323 | |
| 324 | def update(self, minibatches, unlabeled=None): |
| 325 | device = "cuda" if minibatches[0][0].is_cuda else "cpu" |
| 326 | penalty_weight = ( |
| 327 | self.hparams['irm_lambda'] |
| 328 | if self.update_count >= self.hparams['irm_penalty_anneal_iters'] else 1.0 |
| 329 | ) |
| 330 | nll = 0. |
| 331 | penalty = 0. |
| 332 | |
| 333 | all_x = torch.cat([x for x, y in minibatches]) |
| 334 | all_logits = self.network(all_x) |
| 335 | all_logits_idx = 0 |
| 336 | for i, (x, y) in enumerate(minibatches): |
| 337 | logits = all_logits[all_logits_idx:all_logits_idx + x.shape[0]] |
| 338 | all_logits_idx += x.shape[0] |
| 339 | nll += F.cross_entropy(logits, y) |
| 340 | penalty += self._irm_penalty(logits, y) |
| 341 | nll /= len(minibatches) |
| 342 | penalty /= len(minibatches) |
| 343 | loss = nll + (penalty_weight * penalty) |
| 344 | |
| 345 | if self.update_count == self.hparams['irm_penalty_anneal_iters']: |
| 346 | # Reset Adam, because it doesn't like the sharp jump in gradient |
| 347 | # magnitudes that happens at this step. |
| 348 | self.optimizer = torch.optim.Adam( |
| 349 | self.network.parameters(), |
| 350 | lr=self.hparams["lr"], |
| 351 | weight_decay=self.hparams['weight_decay'] |
| 352 | ) |
| 353 | |
| 354 | self.optimizer.zero_grad() |
| 355 | loss.backward() |
| 356 | self.optimizer.step() |
| 357 | |
| 358 | self.update_count += 1 |
| 359 | return {'loss': loss.item(), 'nll': nll.item(), 'penalty': penalty.item()} |
| 360 | |
| 361 | |
| 362 | class VREx(ERM): |
nothing calls this directly
no test coverage detected