def test_correctness(ray_start_2_cpus, num_workers, use_local): layer = nn.Linear(1, 1) ptl_op = TrainingOperator.from_ptl(PTL_Module) trainer1 = TorchTrainer(training_operator_cls=ptl_op, config={ "layer": layer, "data_size": 3, "batch_size": 1 }, num_workers=num_workers, use_local=use_local) train1_stats = trainer1.train() val1_stats = trainer1.validate() trainer1.shutdown() trainer2 = TorchTrainer(training_operator_cls=CorrectnessOperator, scheduler_step_freq="manual", config={ "layer": layer, "data_size": 3, "batch_size": 1 }, num_workers=num_workers, use_local=use_local) train2_stats = trainer2.train() val2_stats = trainer2.validate() trainer2.shutdown() assert train1_stats["train_loss"] == train2_stats["train_loss"] assert val1_stats["val_loss"] == val2_stats["val_loss"] assert val1_stats["val_acc"] == val2_stats["val_accuracy"]
def train_mnist(num_workers=1, use_gpu=False, num_epochs=5): Operator = TrainingOperator.from_ptl(LitMNIST) trainer = TorchTrainer( training_operator_cls=Operator, num_workers=num_workers, config={"lr": 1e-3, "batch_size": 64}, use_gpu=use_gpu, use_tqdm=True, ) for i in range(num_epochs): stats = trainer.train() print(stats) print(trainer.validate()) print("Saving model checkpoint to ./model.pt") trainer.save("./model.pt") print("Model Checkpointed!") trainer.shutdown() print("success!")
self.loss = nn.MSELoss() def train_dataloader(self): return self.train_loader def val_dataloader(self): return self.val_loader def on_save_checkpoint(self, checkpoint): checkpoint["int"] = self.rand_int def on_load_checkpoint(self, checkpoint): self.rand_int = checkpoint["int"] Operator = TrainingOperator.from_ptl(PTL_Module) @pytest.mark.parametrize("use_local", [True, False]) def test_single_step(ray_start_2_cpus, use_local): # noqa: F811 trainer = TorchTrainer(training_operator_cls=Operator, num_workers=1, use_local=use_local, use_gpu=False) metrics = trainer.train(num_steps=1) assert metrics[BATCH_COUNT] == 1 val_metrics = trainer.validate(num_steps=1) assert val_metrics[BATCH_COUNT] == 1 trainer.shutdown()