Progress the pipeline to train the next batch of data. The engine will ingest ``self.train_batch_size()`` total samples collectively across all workers. An iterator that over training data should be provided as an argument unless ``deepspeed.initialize()`` was provided a tr
(self, data_iter=None)
| 339 | self.grad_partition_grad_layer_meta_cache = None |
| 340 | |
| 341 | def train_batch(self, data_iter=None): |
| 342 | """Progress the pipeline to train the next batch of data. The engine will ingest |
| 343 | ``self.train_batch_size()`` total samples collectively across all workers. |
| 344 | |
| 345 | |
| 346 | An iterator that over training data should be provided as an argument |
| 347 | unless ``deepspeed.initialize()`` was provided a training set. In that event, |
| 348 | the training data will automatically be read. |
| 349 | |
| 350 | |
| 351 | .. warning:: |
| 352 | A total of ``self.gradient_accumulation_steps()`` entries will be pulled |
| 353 | from ``data_iter`` by each pipeline. There must be sufficient |
| 354 | data left in ``data_iter`` or else a ``StopIteration`` will halt training. |
| 355 | |
| 356 | DeepSpeed provides a convenience class :class:`deepspeed.utils.RepeatingLoader` |
| 357 | that wraps data loaders to automatically restart upon a ``StopIteration``. |
| 358 | |
| 359 | Args: |
| 360 | data_iter (Iterator, optional): Iterator of training data. |
| 361 | |
| 362 | Returns: |
| 363 | The arithmetic mean of the losses computed this batch. |
| 364 | """ |
| 365 | if not torch._C.is_grad_enabled(): |
| 366 | raise RuntimeError('train_batch() requires gradients enabled. Use eval_batch() instead.') |
| 367 | |
| 368 | # Curriculum learning could change activation shape |
| 369 | if self.curriculum_enabled_legacy(): |
| 370 | new_difficulty = self.curriculum_scheduler_legacy.update_difficulty( \ |
| 371 | self.global_steps + 1) |
| 372 | if self.global_steps == 0 or self.curriculum_scheduler_legacy.first_step: |
| 373 | self.reset_activation_shape() |
| 374 | self.curriculum_scheduler_legacy.first_step = False |
| 375 | elif new_difficulty != self.curriculum_scheduler_legacy.get_difficulty( \ |
| 376 | self.global_steps): |
| 377 | self.reset_activation_shape() |
| 378 | |
| 379 | if data_iter is not None: |
| 380 | self.set_dataiterator(data_iter) |
| 381 | |
| 382 | self.module.train() |
| 383 | self.total_loss = None |
| 384 | self.total_additional_losses = None |
| 385 | self._compute_loss = True |
| 386 | |
| 387 | # Do the work |
| 388 | self.timers(TRAIN_BATCH_TIMER).start() |
| 389 | sched = schedule.TrainSchedule(micro_batches=self.micro_batches, |
| 390 | stages=self.num_stages, |
| 391 | stage_id=self.stage_id) |
| 392 | self._exec_schedule(sched) |
| 393 | |
| 394 | with torch.no_grad(): |
| 395 | self.agg_train_loss = self._aggregate_total_loss() |
| 396 | |
| 397 | self.timers(TRAIN_BATCH_TIMER).stop() |
| 398 |