def test_on_train_end(self): worker = MockWorker() task_data_service = TaskDataService(worker, JobType.TRAINING_WITH_EVALUATION) dataset = tf.data.Dataset.from_tensor_slices( np.array([[1.0, 2.0, 3.0, 4.0], [1.0, 2.0, 3.0, 4.0]])) task_data_service._pending_train_end_callback_task = ( "", 0, 1, elasticdl_pb2.TRAIN_END_CALLBACK, ) task_data_service.get_dataset_by_task = mock.Mock(return_value=dataset) with tempfile.TemporaryDirectory() as temp_dir_name: checkpoint_dir = os.path.join(temp_dir_name, "checkpoint") model = custom_model_with_embedding_layer() save_checkpoint_without_embedding(model, checkpoint_dir) model_handler = ModelHandler.get_model_handler( distribution_strategy=DistributionStrategy.PARAMETER_SERVER, checkpoint_dir=checkpoint_dir, ) saved_model_exporter = SavedModelExporter(task_data_service, dataset_fn, model_handler) saved_model_path = os.path.join(temp_dir_name, "test_exporter") params = {"batch_size": 10, "saved_model_path": saved_model_path} saved_model_exporter.set_params(params) saved_model_exporter.set_model(model) saved_model_exporter.on_train_end() self.assertTrue(os.path.exists(saved_model_path)) self.assertTrue( os.path.exists(os.path.join(saved_model_path, "saved_model.pb")))
def _init_callbacks(self, args): saved_model_exporter = SavedModelExporter(self._task_data_service, self._feed, self._model_handler) # Place default callbacks at the head to execute them firstly self._callbacks_list.callbacks.insert(0, saved_model_exporter) self._callbacks_list.set_model(self._model_inst) set_callback_parameters( self._callbacks_list, batch_size=args.minibatch_size, saved_model_path=args.output, checkpoint_path=args.checkpoint_dir, ) self._saved_model_path = args.output
def _init_from_args(self, args): """ Please refer to elastic/python/common/args.py for more details about arguments of a worker. """ self._worker_id = args.worker_id self._job_type = args.job_type self._minibatch_size = args.minibatch_size self._log_loss_steps = args.log_loss_steps ( model_inst, self._dataset_fn, self._loss, self._opt_fn, self._eval_metrics_fn, self._prediction_outputs_processor, self._custom_data_reader, self._callbacks_list, ) = get_model_spec( model_zoo=args.model_zoo, model_def=args.model_def, dataset_fn=args.dataset_fn, loss=args.loss, optimizer=args.optimizer, eval_metrics_fn=args.eval_metrics_fn, model_params=args.model_params, prediction_outputs_processor=args.prediction_outputs_processor, custom_data_reader=args.custom_data_reader, callbacks=args.callbacks, ) self._collective_communicator = ( CollectiveCommunicator( service_name=args.collective_communicator_service_name ) if self._distribution_strategy == DistributionStrategy.ALLREDUCE else None ) self._model_handler = ModelHandler.get_model_handler( self._distribution_strategy, checkpoint_dir=args.checkpoint_dir ) model_inst = self._model_handler.get_model_to_train(model_inst) self.set_model(model_inst) self._model_version = -1 if self._distribution_strategy != DistributionStrategy.ALLREDUCE: self._model_versions_from_ps = [-1 for _ in range(self._ps_num)] self._task_data_service = TaskDataService( self, self._job_type == JobType.TRAINING_WITH_EVALUATION, data_reader_params=get_dict_from_params_str( args.data_reader_params ), data_origin=args.training_data, ) if self._dataset_fn is None: if hasattr( self._task_data_service.data_reader, "default_dataset_fn" ): self._dataset_fn = ( self._task_data_service.data_reader.default_dataset_fn() ) else: raise ValueError( "dataset_fn is required if the data_reader used does " "not provide default implementation of dataset_fn" ) self._get_model_steps = args.get_model_steps self._opt = self._opt_fn() self._model.optimizer = self._opt self._non_embed_grads = {} self._evaluation_result = {} saved_model_exporter = SavedModelExporter( self._task_data_service, self._dataset_fn, self._model_handler ) # Place default callbacks at the head to execute them firstly self._callbacks_list.callbacks.insert(0, saved_model_exporter) self._callbacks_list.set_model(model_inst) set_callback_parameters( self._callbacks_list, batch_size=args.minibatch_size, saved_model_path=args.output, checkpoint_path=args.checkpoint_dir, )
def _init_from_args(self, args): """ Please refer to elastic/python/common/args.py for more details about arguments of a worker. """ self._worker_id = args.worker_id self._job_type = args.job_type self._minibatch_size = args.minibatch_size self._log_loss_steps = args.log_loss_steps ( model_inst, self._dataset_fn, loss, opt_fn, self._eval_metrics_fn, self._prediction_outputs_processor, self._custom_data_reader, self._callbacks_list, ) = get_model_spec( model_zoo=args.model_zoo, model_def=args.model_def, dataset_fn=args.dataset_fn, loss=args.loss, optimizer=args.optimizer, eval_metrics_fn=args.eval_metrics_fn, prediction_outputs_processor=args.prediction_outputs_processor, custom_data_reader=args.custom_data_reader, callbacks=args.callbacks, ) model_handler = ModelHandler.get_model_handler( self._distribution_strategy, checkpoint_dir=args.checkpoint_dir ) model_inst = model_handler.get_model_to_train(model_inst) model_inst.optimizer = opt_fn() model_inst.loss = loss self._model_version = -1 self._task_data_service = TaskDataService( self._mc, self._job_type == JobType.TRAINING_WITH_EVALUATION, custom_data_reader=self._custom_data_reader, data_reader_params=get_dict_from_params_str( args.data_reader_params ), data_origin=args.training_data, ) if self._dataset_fn is None: if hasattr( self._task_data_service.data_reader, "default_dataset_fn" ): self._dataset_fn = ( self._task_data_service.data_reader.default_dataset_fn() ) else: raise ValueError( "dataset_fn is required if the data_reader used does " "not provide default implementation of dataset_fn" ) self._get_model_steps = args.get_model_steps saved_model_exporter = SavedModelExporter( self._task_data_service, self._dataset_fn, model_handler ) # Place default callbacks at the head to execute them firstly self._callbacks_list.callbacks.insert(0, saved_model_exporter) self._callbacks_list.set_model(model_inst) set_callback_parameters( self._callbacks_list, batch_size=args.minibatch_size, saved_model_path=args.output, checkpoint_path=args.checkpoint_dir, ) self._saved_model_path = args.output if self._distribution_strategy == DistributionStrategy.ALLREDUCE: master_addr = args.master_addr.split(":")[0] self._trainer = AllReduceTrainer(self._mc, master_addr, model_inst) elif ( self._distribution_strategy == DistributionStrategy.PARAMETER_SERVER ): self._trainer = ParameterServerTrainer( model_inst, self._ps_client, self._timing, args )