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__(self, args):
self.logger = get_logger("master", level=args.log_level.upper())
self.num_ps_pods = args.num_ps_pods
self.checkpoint_output_path = args.checkpoint_dir
self.distribution_strategy = args.distribution_strategy
# Master addr
master_ip = os.getenv("MY_POD_IP", "localhost")
self.master_addr = "%s:%d" % (master_ip, args.port)
self.job_type = Master._get_job_type(args)
self.rendezvous_server = None
if self.distribution_strategy == DistributionStrategy.ALLREDUCE:
self.rendezvous_server = HorovodRendezvousServer(master_ip)
# Initialize TensorBoard service if requested
self.tb_service = self._create_tensorboard_service(
args.tensorboard_log_dir, master_ip
)
if self.tb_service:
self.tb_client = TensorBoardClient(
job_name=args.job_name,
image_name=args.worker_image,
namespace=args.namespace,
)
# Initialize the components from the model definition
self.model_module = load_module(
get_module_file_path(args.model_zoo, args.model_def)
).__dict__
self.model_inst = load_model_from_module(
args.model_def, self.model_module, args.model_params
)
self.optimizer = self.model_module[args.optimizer]()
self._create_data_reader_fn = create_data_reader
if args.custom_data_reader in self.model_module:
self._create_data_reader_fn = self.model_module[
args.custom_data_reader
]
# Initialize the callbacks
self.callbacks_list = load_callbacks_from_module(
args.callbacks, self.model_module
)
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._set_completed_steps_by_checkpoint(args.checkpoint_dir_for_init)
# Start task queue
records_per_task = args.minibatch_size * args.num_minibatches_per_task
self.task_d = _make_task_dispatcher(
args.training_data,
args.validation_data,
args.prediction_data,
records_per_task,
args.num_epochs,
args.data_reader_params,
self._create_data_reader_fn,
self.callbacks_list,
)
self.task_d.add_deferred_callback_create_train_end_task()
self.evaluation_service = self._create_evaluation_service(args)
# Initialize instance manager
self.instance_manager = self._create_instance_manager(args)
# Initialize master service
self.master_servicer, self.server = self._create_master_service(args)
self._should_stop = False
self._exit_code = 0
threading.Thread(
target=self._check_timeout_tasks,
name="check_timeout_tasks",
daemon=True,
).start()
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
)