def test_without_init(self): self.optimizer.x = self.init extension = extensions.StepShift('x', self.gamma, self.step, init=self.init, target=self.target) self._run_trainer(extension, self.expect)
def setUp(self): self.optimizer = mock.MagicMock() self.extension = extensions.StepShift('x', self.gamma, self.step, self.init, self.target, self.optimizer) self.interval = 1 self.expect = [e for e in self.expect for _ in range(self.interval)] self.trigger = training.get_trigger((self.interval, 'iteration')) self.trainer = testing.get_trainer_with_mock_updater(self.trigger) self.trainer.updater.get_optimizer.return_value = self.optimizer
def test_resume(self): new_optimizer = mock.Mock() new_extension = extensions.StepShift('x', self.gamma, self.step, self.init, self.target, new_optimizer) self.trainer.extend(self.extension) self.trainer.run() new_trainer = testing.get_trainer_with_mock_updater((5, 'iteration')) new_trainer.extend(new_extension) testing.save_and_load_npz(self.trainer, new_trainer) new_extension.initialize(new_trainer) self.assertEqual(new_optimizer.x, self.optimizer.x) self.assertIsInstance(new_optimizer.x, float)
def test_with_optimizer(self): optimizer = mock.Mock() optimizer.x = 0 extension = extensions.StepShift('x', self.gamma, self.step, self.init, self.target, optimizer) self._run_trainer(extension, self.expect, optimizer)
def create_trainer( config: Config, output: Path, ): assert_config(config) if output.exists(): raise Exception(f"output directory {output} already exists.") # model predictor = create_predictor(config.model) if config.train.trained_model is not None: chainer.serializers.load_npz( config.train.trained_model["predictor_path"], predictor) model = Model( loss_config=config.loss, predictor=predictor, local_padding_size=config.dataset.local_padding_size, ) model.to_gpu(config.train.gpu[0]) cuda.get_device_from_id(config.train.gpu[0]).use() # dataset dataset = create_dataset(config.dataset) batchsize_devided = config.train.batchsize // len(config.train.gpu) train_iter = MultiprocessIterator(dataset["train"], config.train.batchsize) test_iter = MultiprocessIterator(dataset["test"], batchsize_devided, repeat=False, shuffle=True) train_test_iter = MultiprocessIterator(dataset["train_test"], batchsize_devided, repeat=False, shuffle=True) if dataset["test_eval"] is not None: test_eval_iter = MultiprocessIterator(dataset["test_eval"], batchsize_devided, repeat=False, shuffle=True) else: test_eval_iter = None # optimizer def create_optimizer(model): cp: Dict[str, Any] = copy(config.train.optimizer) n = cp.pop("name").lower() if n == "adam": optimizer = optimizers.Adam(**cp) elif n == "sgd": optimizer = optimizers.SGD(**cp) else: raise ValueError(n) optimizer.setup(model) if config.train.optimizer_gradient_clipping is not None: optimizer.add_hook( optimizer_hooks.GradientClipping( config.train.optimizer_gradient_clipping)) return optimizer optimizer = create_optimizer(model) if config.train.trained_model is not None: chainer.serializers.load_npz( config.train.trained_model["optimizer_path"], optimizer) # updater if len(config.train.gpu) <= 1: updater = StandardUpdater( iterator=train_iter, optimizer=optimizer, converter=concat_optional, device=config.train.gpu[0], ) else: updater = ParallelUpdater( iterator=train_iter, optimizer=optimizer, converter=concat_optional, devices={ "main" if i == 0 else f"gpu{gpu}": gpu for i, gpu in enumerate(config.train.gpu) }, ) if config.train.trained_model is not None: updater.iteration = optimizer.t # trainer output.mkdir() config.save_as_json((output / "config.json").absolute()) trigger_log = (config.train.log_iteration, "iteration") trigger_snapshot = (config.train.snapshot_iteration, "iteration") trigger_stop = ((config.train.stop_iteration, "iteration") if config.train.stop_iteration is not None else None) trainer = training.Trainer(updater, stop_trigger=trigger_stop, out=output) tb_writer = SummaryWriter(Path(output)) shift_ext = None if config.train.linear_shift is not None: shift_ext = extensions.LinearShift(**config.train.linear_shift) if config.train.step_shift is not None: shift_ext = extensions.StepShift(**config.train.step_shift) if shift_ext is not None: if config.train.trained_model is not None: shift_ext._t = optimizer.t trainer.extend(shift_ext) if config.train.ema_decay is not None: train_predictor = predictor predictor = deepcopy(predictor) ext = ExponentialMovingAverage(target=train_predictor, ema_target=predictor, decay=config.train.ema_decay) trainer.extend(ext, trigger=(1, "iteration")) ext = extensions.Evaluator(test_iter, model, concat_optional, device=config.train.gpu[0]) trainer.extend(ext, name="test", trigger=trigger_log) ext = extensions.Evaluator(train_test_iter, model, concat_optional, device=config.train.gpu[0]) trainer.extend(ext, name="train", trigger=trigger_log) if test_eval_iter is not None: generator = Generator(config=config, model=predictor, max_batch_size=config.train.batchsize) generate_evaluator = GenerateEvaluator( generator=generator, time_length=config.dataset.time_length_evaluate, local_padding_time_length=config.dataset. local_padding_time_length_evaluate, ) ext = extensions.Evaluator( test_eval_iter, generate_evaluator, concat_optional, device=config.train.gpu[0], ) trainer.extend(ext, name="eval", trigger=trigger_snapshot) ext = extensions.snapshot_object(predictor, filename="main_{.updater.iteration}.npz") trainer.extend(ext, trigger=trigger_snapshot) # ext = extensions.snapshot_object( # optimizer, filename="optimizer_{.updater.iteration}.npz" # ) # trainer.extend(ext, trigger=trigger_snapshot) trainer.extend(extensions.FailOnNonNumber(), trigger=trigger_log) trainer.extend(extensions.observe_lr(), trigger=trigger_log) trainer.extend(extensions.LogReport(trigger=trigger_log)) trainer.extend( extensions.PrintReport(["iteration", "main/loss", "test/main/loss"]), trigger=trigger_log, ) trainer.extend(TensorBoardReport(writer=tb_writer), trigger=trigger_log) trainer.extend(extensions.dump_graph(root_name="main/loss")) if trigger_stop is not None: trainer.extend(extensions.ProgressBar(trigger_stop)) return trainer