def __init__(
self,
config: Config,
model: SILTransformer,
src_spp: Optional[sp.SentencePieceProcessor],
trg_spp: Optional[sp.SentencePieceProcessor],
step: int,
checkpoint_path: Path,
type: str = None,
):
self.types: List[str] = []
if type is not None:
self.types.append(type)
self.step = step
# Configuration priority: user config > auto config > default config.
new_config = copy.deepcopy(_CONFIG_FALLBACK)
merge_dict(new_config, model.auto_config())
merge_dict(new_config, config.root)
new_config["params"].setdefault("num_hypotheses",
new_config["infer"].get("n_best", 1))
new_config["params"].setdefault(
"average_loss_in_time",
new_config["train"]["batch_type"] == "tokens")
new_config["infer"]["n_best"] = 1
self.config = new_config
self.src_spp = src_spp
self.trg_spp = trg_spp
self.model: SILTransformer = clone_layer(model)
self.model.initialize(self.config["data"],
params=self.config["params"])
self._analyze_fn: Optional[Function] = None
self.checkpoint_path = checkpoint_path
self.checkpoint: Checkpoint = None
def _init_model(self, config):
model = misc.clone_layer(self._model)
model.initialize(config["data"], params=config["params"])
if "optimizer" in config["params"]:
optimizer = model.get_optimizer()
else:
optimizer = None
checkpoint = checkpoint_util.Checkpoint(
model,
optimizer=optimizer,
model_dir=config.get("model_dir"),
keep_checkpoint_max=config["train"].get("keep_checkpoint_max", 8))
return checkpoint
def restore(self, checkpoint_path=None, weights_only=False):
"""Restores a checkpoint.
Args:
checkpoint_path: Path a checkpoint to restore. If not set, the latest
checkpoint from :obj:`model_dir` will be restored.
weights_only: Only restore model weights.
Returns:
Path to the restored checkpoint.
"""
if weights_only:
checkpoint = tf.train.Checkpoint(model=self._model)
else:
checkpoint = self._checkpoint
if checkpoint_path is not None:
if tf.io.gfile.isdir(checkpoint_path):
checkpoint_path = tf.train.latest_checkpoint(checkpoint_path)
elif self._checkpoint_manager.latest_checkpoint is not None:
checkpoint_path = self._checkpoint_manager.latest_checkpoint
if checkpoint_path is None:
tf.get_logger().warning("No checkpoint to restore in %s",
self._model_dir)
return None
if is_v1_checkpoint(checkpoint_path):
tf.get_logger().info("Upgrading V1 checkpoint...")
# Work with copies of model and optimizer as the downstream task might
# need to create the variable differently (e.g. under a distribution
# strategy scope).
tmp_model = misc.clone_layer(self._model)
tmp_optimizer = copy.deepcopy(
self._optimizer) if self._optimizer is not None else None
tmp_model.create_variables(optimizer=tmp_optimizer)
step = _restore_v1_checkpoint(checkpoint_path,
tmp_model,
optimizer=tmp_optimizer)
# Save an updated checkpoint in the model directory and restore this one instead.
tmp_checkpoint = Checkpoint(tmp_model,
optimizer=tmp_optimizer,
model_dir=self._model_dir)
checkpoint_path = tmp_checkpoint.save(step)
return self.restore(checkpoint_path=checkpoint_path,
weights_only=weights_only)
load_status = checkpoint.restore(checkpoint_path)
if weights_only:
load_status.expect_partial()
tf.get_logger().info("Restored checkpoint %s", checkpoint_path)
return checkpoint_path
def __init__(self,
model,
config,
auto_config=False,
mixed_precision=False,
seed=None):
"""Initializes the runner parameters.
Args:
model: A :class:`opennmt.models.Model` instance to run or a callable that
returns such instance.
config: The run configuration.
auto_config: If ``True``, use automatic configuration values defined by
:obj:`model`.
mixed_precision: Enable mixed precision.
seed: The random seed to set.
Raises:
TypeError: if :obj:`model` is not a :class:`opennmt.models.Model` instance
or a callable.
"""
if isinstance(model, models.Model):
self._model = model
self._model_fn = lambda: misc.clone_layer(model)
elif callable(model):
self._model = model()
self._model_fn = model
else:
raise TypeError(
"model should be a opennmt.models.Model instance or a callable"
)
tf.get_logger().info("Using OpenNMT-tf version %s", __version__)
tf.get_logger().info("Using model:\n%s", self._model)
self._optimizer = None
self._config = copy.deepcopy(config)
self._auto_config = auto_config
self._mixed_precision = mixed_precision
if mixed_precision:
tf.config.optimizer.set_experimental_options(
{"auto_mixed_precision": True})
if seed is not None:
np.random.seed(seed)
random.seed(seed)
tf.random.set_seed(seed)
def _init_model(self, config):
model = misc.clone_layer(self._model)
model.initialize(config["data"], params=config["params"])
return model
