One training step. Return a dictionary that will be passed to logging.
(
self,
epoch: int,
bidx: int,
batch: T.Any, # xs, ys: (n, dim_x), (n,)
mode: str,
)
| 162 | self.test_sampler.set_epoch(epoch) |
| 163 | |
| 164 | def _step( |
| 165 | self, |
| 166 | epoch: int, |
| 167 | bidx: int, |
| 168 | batch: T.Any, # xs, ys: (n, dim_x), (n,) |
| 169 | mode: str, |
| 170 | ) -> T.Dict[str, T.Any]: |
| 171 | """One training step. |
| 172 | Return a dictionary that will be passed to logging. |
| 173 | """ |
| 174 | xs, ys_gt = batch # (n, dim_x), (n,) |
| 175 | xs = xs.to(self.device) |
| 176 | ys_gt = ys_gt.to(self.device) |
| 177 | ys = self.net(xs) |
| 178 | loss = (ys_gt - ys.squeeze(-1)).pow(2).mean() |
| 179 | |
| 180 | if mode == 'train': |
| 181 | # step optimizer |
| 182 | self.optimizer.zero_grad() |
| 183 | loss.backward() |
| 184 | self.optimizer.step() |
| 185 | else: |
| 186 | pass |
| 187 | |
| 188 | if mode not in self.loss_values: |
| 189 | self.loss_values[mode] = [] |
| 190 | self.step_values[mode] = [] |
| 191 | |
| 192 | if mode == 'train': |
| 193 | self.step_count += 1 |
| 194 | |
| 195 | self.loss_values[mode].append(loss.detach().cpu().item()) |
| 196 | self.step_values[mode].append(self.step_count) |
| 197 | |
| 198 | return dict( |
| 199 | loss=loss.detach().cpu().item() |
| 200 | ) |
| 201 | |
| 202 | @customer |
| 203 | def train_step( |
no test coverage detected