def evaluate_gridworld(self, environment, evaluator, trainer, use_gpu, one_hot_action=True): pre_training_loss = None if self.run_pre_training_eval: predictor = self.get_predictor(trainer, environment) evaluator.evaluate(predictor) print("Pre-Training eval: ", evaluator.mc_loss[-1]) pre_training_loss = evaluator.mc_loss[-1] for _ in range(self.num_epochs): samples = environment.generate_samples(10240, 1.0, DISCOUNT) if trainer.rl_parameters.reward_boost is not None: # Reverse any reward boost rewards_update = [] for action, reward in zip(samples.actions, samples.rewards): rewards_update.append( reward - trainer.rl_parameters.reward_boost.get(action, 0.0)) samples = samples._replace(rewards=rewards_update) tdps = environment.preprocess_samples( samples, self.minibatch_size, use_gpu=use_gpu, one_hot_action=one_hot_action, ) for tdp in tdps: trainer.train(tdp) predictor = self.get_predictor(trainer, environment) evaluator.evaluate(predictor) print("Post-Training eval: ", evaluator.mc_loss[-1]) if self.check_tolerance: self.assertLess(evaluator.mc_loss[-1], self.tolerance_threshold) if self.run_pre_training_eval: self.assertGreater(pre_training_loss, evaluator.mc_loss[-1]) if self.test_save_load: with tempfile.TemporaryDirectory() as tmpdirname: tmp_path = os.path.join(tmpdirname, "model") serving_module_bytes = export_module_to_buffer( predictor.model).getvalue() logger.info( "Saving TorchScript predictor to {}".format(tmp_path)) with open(tmp_path, "wb") as output_fp: output_fp.write(serving_module_bytes) new_predictor = type(predictor)(torch.jit.load(tmp_path)) evaluator.evaluate(new_predictor) print("Post-ONNX eval: ", evaluator.mc_loss[-1]) if self.check_tolerance: self.assertLess(evaluator.mc_loss[-1], self.tolerance_threshold)
def test_predictor_torch_export(self): """Verify that q-values before model export equal q-values after model export. Meant to catch issues with export logic.""" environment = Gridworld() samples = Samples( mdp_ids=["0"], sequence_numbers=[0], sequence_number_ordinals=[1], states=[{0: 1.0, 1: 1.0, 2: 1.0, 3: 1.0, 4: 1.0, 5: 1.0, 15: 1.0, 24: 1.0}], actions=["D"], action_probabilities=[0.5], rewards=[0], possible_actions=[["R", "D"]], next_states=[{5: 1.0}], next_actions=["U"], terminals=[False], possible_next_actions=[["R", "U", "D"]], ) tdps = environment.preprocess_samples(samples, 1) assert len(tdps) == 1, "Invalid number of data pages" trainer, exporter = self.get_modular_sarsa_trainer_exporter( environment, {}, False ) input = rlt.PreprocessedState.from_tensor(tdps[0].states) pre_export_q_values = trainer.q_network(input).q_values.detach().numpy() preprocessor = Preprocessor(environment.normalization, False) cpu_q_network = trainer.q_network.cpu_model() cpu_q_network.eval() dqn_with_preprocessor = DiscreteDqnWithPreprocessor(cpu_q_network, preprocessor) serving_module = DiscreteDqnPredictorWrapper( dqn_with_preprocessor, action_names=environment.ACTIONS ) with tempfile.TemporaryDirectory() as tmpdirname: buf = export_module_to_buffer(serving_module) tmp_path = os.path.join(tmpdirname, "model") with open(tmp_path, "wb") as f: f.write(buf.getvalue()) f.close() predictor = DiscreteDqnTorchPredictor(torch.jit.load(tmp_path)) post_export_q_values = predictor.predict([samples.states[0]]) for i, action in enumerate(environment.ACTIONS): self.assertAlmostEqual( float(pre_export_q_values[0][i]), float(post_export_q_values[0][action]), places=4, )
def save_torchscript_model(self, module: torch.nn.Module, path: str): logger.info("Saving TorchScript predictor to {}".format(path)) serving_module_bytes = export_module_to_buffer(module).getvalue() with open(path, "wb") as output_fp: output_fp.write(serving_module_bytes)