A simple trainer for the most common type of task: single-cost single-optimizer single-data-source iterative optimization, optionally using data-parallelism. It assumes that every step, you: 1. Compute the loss with a data from the data_loader. 2. Compute the gradients with
| 169 | |
| 170 | |
| 171 | class SimpleTrainer(TrainerBase): |
| 172 | """ |
| 173 | A simple trainer for the most common type of task: |
| 174 | single-cost single-optimizer single-data-source iterative optimization, |
| 175 | optionally using data-parallelism. |
| 176 | It assumes that every step, you: |
| 177 | |
| 178 | 1. Compute the loss with a data from the data_loader. |
| 179 | 2. Compute the gradients with the above loss. |
| 180 | 3. Update the model with the optimizer. |
| 181 | |
| 182 | All other tasks during training (checkpointing, logging, evaluation, LR schedule) |
| 183 | are maintained by hooks, which can be registered by :meth:`TrainerBase.register_hooks`. |
| 184 | |
| 185 | If you want to do anything fancier than this, |
| 186 | either subclass TrainerBase and implement your own `run_step`, |
| 187 | or write your own training loop. |
| 188 | """ |
| 189 | |
| 190 | def __init__(self, model, data_loader, optimizer): |
| 191 | """ |
| 192 | Args: |
| 193 | model: a torch Module. Takes a data from data_loader and returns a |
| 194 | dict of losses. |
| 195 | data_loader: an iterable. Contains data to be used to call model. |
| 196 | optimizer: a torch optimizer. |
| 197 | """ |
| 198 | super().__init__() |
| 199 | |
| 200 | """ |
| 201 | We set the model to training mode in the trainer. |
| 202 | However it's valid to train a model that's in eval mode. |
| 203 | If you want your model (or a submodule of it) to behave |
| 204 | like evaluation during training, you can overwrite its train() method. |
| 205 | """ |
| 206 | model.train() |
| 207 | |
| 208 | self.model = model |
| 209 | self.data_loader = data_loader |
| 210 | self._data_loader_iter = iter(data_loader) |
| 211 | self.optimizer = optimizer |
| 212 | |
| 213 | def run_step(self): |
| 214 | """ |
| 215 | Implement the standard training logic described above. |
| 216 | """ |
| 217 | assert self.model.training, "[SimpleTrainer] model was changed to eval mode!" |
| 218 | start = time.perf_counter() |
| 219 | """ |
| 220 | If you want to do something with the data, you can wrap the dataloader. |
| 221 | """ |
| 222 | data = next(self._data_loader_iter) |
| 223 | data_time = time.perf_counter() - start |
| 224 | |
| 225 | """ |
| 226 | If you want to do something with the losses, you can wrap the model. |
| 227 | """ |
| 228 | loss_dict = self.model(data) |