def train( self, train_dataset: Optional[Dataset], eval_dataset: Optional[Dataset], data_module: Optional[ReAgentDataModule], num_epochs: int, reader_options: ReaderOptions, resource_options: Optional[ResourceOptions], ) -> RLTrainingOutput: batch_preprocessor = self.build_batch_preprocessor() reporter = self.get_reporter() # pyre-fixme[16]: `Trainer` has no attribute `set_reporter`. # pyre-fixme[16]: `Trainer` has no attribute `set_reporter`. self.trainer.set_reporter(reporter) # assert eval_dataset is None self._lightning_trainer = train_eval_lightning( train_dataset=train_dataset, eval_dataset=eval_dataset, trainer_module=self.trainer, data_module=data_module, num_epochs=num_epochs, use_gpu=self.use_gpu, batch_preprocessor=batch_preprocessor, reader_options=self.reader_options, checkpoint_path=self._lightning_checkpoint_path, resource_options=resource_options or ResourceOptions(), ) # pyre-fixme[16]: `RLTrainingReport` has no attribute `make_union_instance`. training_report = RLTrainingReport.make_union_instance( reporter.generate_training_report()) return RLTrainingOutput(training_report=training_report)
def train( self, train_dataset: Optional[Dataset], eval_dataset: Optional[Dataset], test_dataset: Optional[Dataset], data_module: Optional[ReAgentDataModule], num_epochs: int, reader_options: ReaderOptions, resource_options: ResourceOptions, ) -> RLTrainingOutput: batch_preprocessor = self.build_batch_preprocessor( resource_options.use_gpu) reporter = self.get_reporter() # pyre-fixme[16]: `Trainer` has no attribute `set_reporter`. # pyre-fixme[16]: `Trainer` has no attribute `set_reporter`. self.trainer.set_reporter(reporter) # assert eval_dataset is None # pyre-fixme[16]: `ActorCriticBase` has no attribute `_lightning_trainer`. self._lightning_trainer = train_eval_lightning( train_dataset=train_dataset, eval_dataset=eval_dataset, test_dataset=test_dataset, trainer_module=self.trainer, data_module=data_module, num_epochs=num_epochs, logger_name="ActorCritic", batch_preprocessor=batch_preprocessor, reader_options=self.reader_options, checkpoint_path=self._lightning_checkpoint_path, resource_options=resource_options or ResourceOptions(), ) if reporter is None: training_report = None else: # pyre-fixme[16]: `RLTrainingReport` has no attribute `make_union_instance`. training_report = RLTrainingReport.make_union_instance( reporter.generate_training_report()) logger_data = self._lightning_trainer.logger.line_plot_aggregated self._lightning_trainer.logger.clear_local_data() return RLTrainingOutput(training_report=training_report, logger_data=logger_data)
def __init__( self, *, input_table_spec: Optional[TableSpec] = None, reward_options: Optional[RewardOptions] = None, setup_data: Optional[Dict[str, bytes]] = None, saved_setup_data: Optional[Dict[str, bytes]] = None, reader_options: Optional[ReaderOptions] = None, resource_options: Optional[ResourceOptions] = None, model_manager=None, ): super().__init__() self.input_table_spec = input_table_spec self.reward_options = reward_options or RewardOptions() self.reader_options = reader_options or ReaderOptions() self.resource_options = resource_options or ResourceOptions(gpu=0) self._model_manager = model_manager self.setup_data = setup_data self.saved_setup_data = saved_setup_data or {} self._setup_done = False
def query_and_train( input_table_spec: TableSpec, model: ModelManager__Union, num_epochs: int, use_gpu: bool, *, setup_data: Optional[Dict[str, bytes]] = None, saved_setup_data: Optional[Dict[str, bytes]] = None, normalization_data_map: Optional[Dict[str, NormalizationData]] = None, reward_options: Optional[RewardOptions] = None, reader_options: Optional[ReaderOptions] = None, resource_options: Optional[ResourceOptions] = None, warmstart_path: Optional[str] = None, validator: Optional[ModelValidator__Union] = None, publisher: Optional[ModelPublisher__Union] = None, named_model_ids: Optional[ModuleNameToEntityId] = None, recurring_period: Optional[RecurringPeriod] = None, ) -> RLTrainingOutput: child_workflow_id = get_workflow_id() if named_model_ids is None: named_model_ids = get_new_named_entity_ids(model.value.serving_module_names()) logger.info("Starting query") reward_options = reward_options or RewardOptions() reader_options = reader_options or ReaderOptions() resource_options = resource_options or ResourceOptions() manager = model.value if saved_setup_data is not None: def _maybe_get_bytes(v) -> bytes: if isinstance(v, bytes): return v # HACK: FBLearner sometimes pack bytes into Blob return v.data saved_setup_data = {k: _maybe_get_bytes(v) for k, v in saved_setup_data.items()} if setup_data is None: data_module = manager.get_data_module( input_table_spec=input_table_spec, reward_options=reward_options, reader_options=reader_options, saved_setup_data=saved_setup_data, ) if data_module is not None: setup_data = data_module.prepare_data() # Throw away existing normalization data map normalization_data_map = None if sum([int(setup_data is not None), int(normalization_data_map is not None)]) != 1: raise ValueError("setup_data and normalization_data_map are mutually exclusive") train_dataset = None eval_dataset = None if normalization_data_map is not None: calc_cpe_in_training = manager.should_generate_eval_dataset sample_range_output = get_sample_range(input_table_spec, calc_cpe_in_training) train_dataset = manager.query_data( input_table_spec=input_table_spec, sample_range=sample_range_output.train_sample_range, reward_options=reward_options, ) eval_dataset = None if calc_cpe_in_training: eval_dataset = manager.query_data( input_table_spec=input_table_spec, sample_range=sample_range_output.eval_sample_range, reward_options=reward_options, ) logger.info("Starting training") results = manager.train_workflow( train_dataset, eval_dataset, num_epochs=num_epochs, use_gpu=use_gpu, setup_data=setup_data, normalization_data_map=normalization_data_map, named_model_ids=named_model_ids, child_workflow_id=child_workflow_id, reward_options=reward_options, reader_options=reader_options, resource_options=resource_options, warmstart_path=warmstart_path, ) if validator is not None: results = run_validator(validator, results) if publisher is not None: results = run_publisher( publisher, model, results, named_model_ids, child_workflow_id, recurring_period, ) return results
def train_workflow( model_manager: ModelManager, train_dataset: Optional[Dataset], eval_dataset: Optional[Dataset], *, num_epochs: int, use_gpu: bool, named_model_ids: ModuleNameToEntityId, child_workflow_id: int, setup_data: Optional[Dict[str, bytes]] = None, normalization_data_map: Optional[Dict[str, NormalizationData]] = None, reward_options: Optional[RewardOptions] = None, reader_options: Optional[ReaderOptions] = None, resource_options: Optional[ResourceOptions] = None, warmstart_path: Optional[str] = None, ) -> RLTrainingOutput: writer = SummaryWriter() logger.info("TensorBoard logging location is: {}".format(writer.log_dir)) if setup_data is not None: data_module = model_manager.get_data_module( setup_data=setup_data, reward_options=reward_options, reader_options=reader_options, resource_options=resource_options, ) assert data_module is not None data_module.setup() else: data_module = None if normalization_data_map is None: assert data_module is not None normalization_data_map = data_module.get_normalization_data_map() warmstart_input_path = warmstart_path or None trainer_module = model_manager.build_trainer( use_gpu=use_gpu, reward_options=reward_options, normalization_data_map=normalization_data_map, ) if not reader_options: reader_options = ReaderOptions() if not resource_options: resource_options = ResourceOptions() with summary_writer_context(writer): train_output, lightning_trainer = model_manager.train( trainer_module, train_dataset, eval_dataset, None, data_module, num_epochs, reader_options, resource_options, checkpoint_path=warmstart_input_path, ) output_paths = {} for module_name, serving_module in model_manager.build_serving_modules( trainer_module, normalization_data_map).items(): torchscript_output_path = f"{model_manager.__class__.__name__}_{module_name}_{round(time.time())}.torchscript" torch.jit.save(serving_module, torchscript_output_path) logger.info(f"Saved {module_name} to {torchscript_output_path}") output_paths[module_name] = torchscript_output_path return dataclasses.replace(train_output, output_paths=output_paths)
def query_and_train( input_table_spec: TableSpec, model: ModelManager__Union, normalization_data_map: Dict[str, NormalizationData], num_epochs: int, use_gpu: bool, reward_options: Optional[RewardOptions] = None, reader_options: Optional[ReaderOptions] = None, resource_options: Optional[ResourceOptions] = None, warmstart_path: Optional[str] = None, validator: Optional[ModelValidator__Union] = None, publisher: Optional[ModelPublisher__Union] = None, parent_workflow_id: Optional[int] = None, recurring_period: Optional[RecurringPeriod] = None, ) -> RLTrainingOutput: child_workflow_id = get_workflow_id() if parent_workflow_id is None: parent_workflow_id = child_workflow_id logger.info("Starting query") reward_options = reward_options or RewardOptions() reader_options = reader_options or ReaderOptions() resource_options = resource_options or ResourceOptions() manager = model.value calc_cpe_in_training = manager.should_generate_eval_dataset sample_range_output = get_sample_range(input_table_spec, calc_cpe_in_training) train_dataset = manager.query_data( input_table_spec=input_table_spec, sample_range=sample_range_output.train_sample_range, reward_options=reward_options, ) eval_dataset = None if calc_cpe_in_training: eval_dataset = manager.query_data( input_table_spec=input_table_spec, sample_range=sample_range_output.eval_sample_range, reward_options=reward_options, ) logger.info("Starting training") results = manager.train_workflow( train_dataset, eval_dataset, normalization_data_map, num_epochs, use_gpu, parent_workflow_id=parent_workflow_id, child_workflow_id=child_workflow_id, reward_options=reward_options, reader_options=reader_options, resource_options=resource_options, warmstart_path=warmstart_path, ) if validator is not None: results = run_validator(validator, results) if publisher is not None: results = run_publisher( publisher, model, results, parent_workflow_id, child_workflow_id, recurring_period, ) return results
def train_workflow( self, train_dataset: Optional[Dataset], eval_dataset: Optional[Dataset], *, num_epochs: int, use_gpu: bool, named_model_ids: ModuleNameToEntityId, child_workflow_id: int, setup_data: Optional[Dict[str, bytes]] = None, normalization_data_map: Optional[Dict[str, NormalizationData]] = None, reward_options: Optional[RewardOptions] = None, reader_options: Optional[ReaderOptions] = None, resource_options: Optional[ResourceOptions] = None, warmstart_path: Optional[str] = None, ) -> RLTrainingOutput: writer = SummaryWriter() logger.info("TensorBoard logging location is: {}".format( writer.log_dir)) if setup_data is not None: data_module = self.get_data_module(setup_data=setup_data, reader_options=reader_options) assert data_module is not None data_module.setup() else: data_module = None if normalization_data_map is None: assert data_module is not None normalization_data_map = data_module.get_normalization_data_map( self.required_normalization_keys) warmstart_input_path = warmstart_path or None self.initialize_trainer( use_gpu=use_gpu, # pyre-fixme[6]: Expected `RewardOptions` for 2nd param but got # `Optional[RewardOptions]`. reward_options=reward_options, normalization_data_map=normalization_data_map, warmstart_path=warmstart_input_path, ) if not reader_options: reader_options = ReaderOptions() if not resource_options: resource_options = ResourceOptions() with summary_writer_context(writer): train_output = self.train( train_dataset, eval_dataset, data_module, num_epochs, reader_options, resource_options, ) output_paths = {} for module_name, serving_module in self.build_serving_modules().items( ): # TODO: make this a parameter torchscript_output_path = f"model_{round(time.time())}.torchscript" serving_module = self.build_serving_module() torch.jit.save(serving_module, torchscript_output_path) logger.info(f"Saved {module_name} to {torchscript_output_path}") output_paths[module_name] = torchscript_output_path return dataclasses.replace(train_output, output_paths=output_paths)