| 75 | |
| 76 | class Trainer: |
| 77 | def __init__(self, logger, config): |
| 78 | if torch.cuda.is_available(): |
| 79 | self.device = torch.device('cuda') |
| 80 | else: |
| 81 | self.device = torch.device('cpu') |
| 82 | |
| 83 | self.logger = logger |
| 84 | self.train_config = registry.instantiate(TrainConfig, config['train']) |
| 85 | self.data_random = random_state.RandomContext(self.train_config.data_seed) |
| 86 | self.model_random = random_state.RandomContext(self.train_config.model_seed) |
| 87 | |
| 88 | self.init_random = random_state.RandomContext(self.train_config.init_seed) |
| 89 | with self.init_random: |
| 90 | # 0. Construct preprocessors |
| 91 | self.model_preproc = registry.instantiate( |
| 92 | registry.lookup('model', config['model']).Preproc, |
| 93 | config['model'], |
| 94 | unused_keys=('name',)) |
| 95 | self.model_preproc.load() |
| 96 | |
| 97 | # 1. Construct model |
| 98 | self.model = registry.construct('model', config['model'], |
| 99 | unused_keys=('encoder_preproc', 'decoder_preproc'), preproc=self.model_preproc, device=self.device) |
| 100 | self.model.to(self.device) |
| 101 | |
| 102 | def train(self, config, modeldir): |
| 103 | # slight difference here vs. unrefactored train: The init_random starts over here. Could be fixed if it was important by saving random state at end of init |