(self, minibatches, unlabeled=False)
| 1200 | ) |
| 1201 | |
| 1202 | def update(self, minibatches, unlabeled=False): |
| 1203 | assert len(minibatches) == self.num_domains |
| 1204 | all_x = torch.cat([x for x, y in minibatches]) |
| 1205 | all_y = torch.cat([y for x, y in minibatches]) |
| 1206 | len_minibatches = [x.shape[0] for x, y in minibatches] |
| 1207 | |
| 1208 | all_z = self.featurizer(all_x) |
| 1209 | all_logits = self.classifier(all_z) |
| 1210 | |
| 1211 | penalty = self.compute_fishr_penalty(all_logits, all_y, len_minibatches) |
| 1212 | all_nll = F.cross_entropy(all_logits, all_y) |
| 1213 | |
| 1214 | penalty_weight = 0 |
| 1215 | if self.update_count >= self.hparams["penalty_anneal_iters"]: |
| 1216 | penalty_weight = self.hparams["lambda"] |
| 1217 | if self.update_count == self.hparams["penalty_anneal_iters"] != 0: |
| 1218 | # Reset Adam as in IRM or V-REx, because it may not like the sharp jump in |
| 1219 | # gradient magnitudes that happens at this step. |
| 1220 | self._init_optimizer() |
| 1221 | self.update_count += 1 |
| 1222 | |
| 1223 | objective = all_nll + penalty_weight * penalty |
| 1224 | self.optimizer.zero_grad() |
| 1225 | objective.backward() |
| 1226 | self.optimizer.step() |
| 1227 | |
| 1228 | return {'loss': objective.item(), 'nll': all_nll.item(), 'penalty': penalty.item()} |
| 1229 | |
| 1230 | def compute_fishr_penalty(self, all_logits, all_y, len_minibatches): |
| 1231 | dict_grads = self._get_grads(all_logits, all_y) |
nothing calls this directly
no test coverage detected