Class that controls the training process. Args: model(:py:class:`onmt.models.models.NMTModel`): translation models to train train_loss(:obj:`onmt.utils.loss.LossComputeBase`): training loss computation valid_loss(:obj:`onmt
| 61 | |
| 62 | |
| 63 | class Trainer(object): |
| 64 | """ |
| 65 | Class that controls the training process. |
| 66 | |
| 67 | Args: |
| 68 | model(:py:class:`onmt.models.models.NMTModel`): translation models |
| 69 | to train |
| 70 | train_loss(:obj:`onmt.utils.loss.LossComputeBase`): |
| 71 | training loss computation |
| 72 | valid_loss(:obj:`onmt.utils.loss.LossComputeBase`): |
| 73 | training loss computation |
| 74 | optimizer(:obj:`onmt.utils.optimizers.Optimizer`): |
| 75 | the optimizer responsible for update |
| 76 | trunc_size(int): length of truncated back propagation through time |
| 77 | shard_size(int): compute loss in shards of this size for efficiency |
| 78 | data_type(string): type of the source input: [text|img|audio] |
| 79 | norm_method(string): normalization methods: [sents|tokens] |
| 80 | grad_accum_count(int): accumulate gradients this many times. |
| 81 | report_manager(:obj:`onmt.utils.ReportMgrBase`): |
| 82 | the object that creates reports, or None |
| 83 | model_saver(:obj:`onmt.models.ModelSaverBase`): the saver is |
| 84 | used to save a checkpoint. |
| 85 | Thus nothing will be saved if this parameter is None |
| 86 | """ |
| 87 | |
| 88 | def __init__(self, args, model, optimizer, grad_accum_count=1, n_gpu=1, gpu_rank=1, report_manager=None): |
| 89 | # Basic attributes. |
| 90 | self.args = args |
| 91 | self.save_checkpoint_steps = args.save_checkpoint_steps |
| 92 | self.model = model |
| 93 | self.optimizer = optimizer |
| 94 | self.grad_accum_count = grad_accum_count |
| 95 | self.n_gpu = n_gpu |
| 96 | self.gpu_rank = gpu_rank |
| 97 | self.report_manager = report_manager |
| 98 | |
| 99 | self.loss = torch.nn.BCELoss(reduction='none') |
| 100 | assert grad_accum_count > 0 |
| 101 | # Set models in training mode. |
| 102 | if model: |
| 103 | self.model.train() |
| 104 | |
| 105 | def train(self, train_iter_fct, train_steps, valid_iter_fct=None, valid_steps=-1): |
| 106 | """ |
| 107 | The main training loops. |
| 108 | by iterating over training data (i.e. `train_iter_fct`) |
| 109 | and running validation (i.e. iterating over `valid_iter_fct` |
| 110 | |
| 111 | Args: |
| 112 | train_iter_fct(function): a function that returns the train |
| 113 | iterator. e.g. something like |
| 114 | train_iter_fct = lambda: generator(*args, **kwargs) |
| 115 | valid_iter_fct(function): same as train_iter_fct, for valid data |
| 116 | train_steps(int): |
| 117 | valid_steps(int): |
| 118 | save_checkpoint_steps(int): |
| 119 | |
| 120 | Return: |