Args: model: a torch Module. Takes a data from data_loader and returns a dict of losses. data_loader: an iterable. Contains data to be used to call model. optimizer: a torch optimizer.
(self, model, data_loader, optimizer)
| 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 | """ |