def test_trainer_from_base_class_params(self):
params = Params.from_file(self.FIXTURES_ROOT / 'simple_tagger' / 'experiment.json')
# Can instantiate from base class params
TrainerBase.from_params(params, self.TEST_DIR)
def train_model(params: Params,
serialization_dir: str,
file_friendly_logging: bool = False,
recover: bool = False,
force: bool = False,
cache_directory: str = None,
cache_prefix: str = None) -> Model:
"""
Trains the model specified in the given :class:`Params` object, using the data and training
parameters also specified in that object, and saves the results in ``serialization_dir``.
Parameters
----------
params : ``Params``
A parameter object specifying an AllenNLP Experiment.
serialization_dir : ``str``
The directory in which to save results and logs.
file_friendly_logging : ``bool``, optional (default=False)
If ``True``, we add newlines to tqdm output, even on an interactive terminal, and we slow
down tqdm's output to only once every 10 seconds.
recover : ``bool``, optional (default=False)
If ``True``, we will try to recover a training run from an existing serialization
directory. This is only intended for use when something actually crashed during the middle
of a run. For continuing training a model on new data, see the ``fine-tune`` command.
force : ``bool``, optional (default=False)
If ``True``, we will overwrite the serialization directory if it already exists.
cache_directory : ``str``, optional
For caching data pre-processing. See :func:`allennlp.training.util.datasets_from_params`.
cache_prefix : ``str``, optional
For caching data pre-processing. See :func:`allennlp.training.util.datasets_from_params`.
Returns
-------
best_model: ``Model``
The model with the best epoch weights.
"""
create_serialization_dir(params, serialization_dir, recover, force)
stdout_handler = prepare_global_logging(serialization_dir, file_friendly_logging)
prepare_environment(params)
cuda_device = params.params.get('trainer').get('cuda_device', -1)
check_for_gpu(cuda_device)
params.to_file(os.path.join(serialization_dir, CONFIG_NAME))
evaluate_on_test = params.pop_bool("evaluate_on_test", False)
trainer_type = params.get("trainer", {}).get("type", "default")
if trainer_type == "default":
# Special logic to instantiate backward-compatible trainer.
pieces = TrainerPieces.from_params(params, # pylint: disable=no-member
serialization_dir,
recover,
cache_directory,
cache_prefix)
trainer = Trainer.from_params(
model=pieces.model,
serialization_dir=serialization_dir,
iterator=pieces.iterator,
train_data=pieces.train_dataset,
validation_data=pieces.validation_dataset,
params=pieces.params,
validation_iterator=pieces.validation_iterator)
evaluation_iterator = pieces.validation_iterator or pieces.iterator
evaluation_dataset = pieces.test_dataset
else:
if evaluate_on_test:
raise ValueError("--evaluate-on-test only works with the default Trainer. "
"If you're using the CallbackTrainer you can use a callback "
"to evaluate at Events.TRAINING_END; otherwise you'll have "
"to run allennlp evaluate separately.")
trainer = TrainerBase.from_params(params, serialization_dir, recover)
evaluation_dataset = None
params.assert_empty('base train command')
try:
metrics = trainer.train()
except KeyboardInterrupt:
# if we have completed an epoch, try to create a model archive.
if os.path.exists(os.path.join(serialization_dir, _DEFAULT_WEIGHTS)):
logging.info("Training interrupted by the user. Attempting to create "
"a model archive using the current best epoch weights.")
archive_model(serialization_dir, files_to_archive=params.files_to_archive)
raise
# Evaluate
if evaluation_dataset and evaluate_on_test:
logger.info("The model will be evaluated using the best epoch weights.")
test_metrics = evaluate(trainer.model, evaluation_dataset, evaluation_iterator,
cuda_device=trainer._cuda_devices[0], # pylint: disable=protected-access,
# TODO(brendanr): Pass in an arg following Joel's trainer refactor.
batch_weight_key="")
for key, value in test_metrics.items():
metrics["test_" + key] = value
elif evaluation_dataset:
logger.info("To evaluate on the test set after training, pass the "
"'evaluate_on_test' flag, or use the 'allennlp evaluate' command.")
cleanup_global_logging(stdout_handler)
# Now tar up results
archive_model(serialization_dir, files_to_archive=params.files_to_archive)
dump_metrics(os.path.join(serialization_dir, "metrics.json"), metrics, log=True)
# We count on the trainer to have the model with best weights
return trainer.model
}
},
"num_epochs":
num_epochs,
"callbacks":
["generate_training_batches", "train-gan", "track_metrics"]
}
})
class GanCallbackTrainerTest(ModelTestCase):
def test_gan_can_train(self):
params = config(batches_per_epoch=2, num_epochs=2)
train_model(params, self.TEST_DIR)
if __name__ == "__main__":
# Run it yourself, it's fun!
#
# python -m allennlp.tests.training.gan_callback_trainer_test
#
# pylint: disable=invalid-name
from allennlp.training.trainer_base import TrainerBase
serialization_dir = tempfile.mkdtemp()
params = config()
trainer = TrainerBase.from_params(params, serialization_dir)
metrics = trainer.train()
print(metrics)
def _train_worker(
process_rank: int,
params: Params,
serialization_dir: str,
file_friendly_logging: bool = False,
recover: bool = False,
include_package: List[str] = None,
batch_weight_key: str = "",
node_rank: int = 0,
master_addr: str = "127.0.0.1",
master_port: int = 29500,
world_size: int = 1,
distributed_device_ids: List[str] = None,
# For fine-tuning:
model: Model = None,
extend_vocab: bool = False,
embedding_sources_mapping: Dict[str, str] = None,
) -> Optional[Model]:
"""
Helper to train the configured model/experiment. In distributed mode, this is spawned as a
worker process. In a single GPU experiment, this returns the ``Model`` object and in distributed
training, nothing is returned.
# Parameters
process_rank : ``int``
The process index that is initialized using the GPU device id.
params : ``Params``
A parameter object specifying an AllenNLP Experiment.
serialization_dir : ``str``
The directory in which to save results and logs.
file_friendly_logging : ``bool``, optional (default=False)
If ``True``, we add newlines to tqdm output, even on an interactive terminal, and we slow
down tqdm's output to only once every 10 seconds.
recover : ``bool``, optional (default=False)
If ``True``, we will try to recover a training run from an existing serialization
directory. This is only intended for use when something actually crashed during the middle
of a run. For continuing training a model on new data, see the ``fine-tune`` command.
include_package : ``List[str]``, optional
In distributed mode, since this function would have been spawned as a separate process,
the extra imports need to be done again. NOTE: This does not have any effect in single
GPU training.
node_rank : ``int``, optional
Rank of the node
world_size : ``int``, optional
The number of processes involved in distributed training.
# Returns
best_model : ``Model``
The model with the best epoch weights.
"""
prepare_global_logging(serialization_dir,
file_friendly_logging,
rank=process_rank,
world_size=world_size)
prepare_environment(params)
distributed = world_size > 1
# not using `allennlp.common.util.is_master` as the process group is yet to be initialized
master = process_rank == 0
evaluate_on_test = params.pop_bool("evaluate_on_test", False)
if distributed:
# Since the worker is spawned and not forked, the extra imports
# need to be done again.
if include_package is not None:
for package_name in include_package:
import_submodules(package_name)
num_procs_per_node = len(distributed_device_ids)
# The Unique identifier of the worker process among all the processes in the
# distributed training group is computed here. This is used while initializing
# the process group using `init_process_group`
global_rank = node_rank * num_procs_per_node + process_rank
# Number of processes per node is useful to know if a process
# is a master in the local node(node in which it is running)
os.environ["ALLENNLP_PROCS_PER_NODE"] = str(num_procs_per_node)
# In distributed training, the configured device is always going to be a list.
# The corresponding gpu id for the particular worker is obtained by picking the id
# from the device list with the rank as index
gpu_id = distributed_device_ids[process_rank] # type: ignore
# Till now, "cuda_device" might not be set in the trainer params.
# But a worker trainer needs to only know about its specific GPU id.
params["trainer"]["cuda_device"] = gpu_id
params["trainer"]["world_size"] = world_size
params["trainer"]["distributed"] = True
torch.cuda.set_device(int(gpu_id))
dist.init_process_group(
backend="nccl",
init_method=f"tcp://{master_addr}:{master_port}",
world_size=world_size,
rank=global_rank,
)
logging.info(f"Process group of world size {world_size} initialized "
f"for distributed training in worker {global_rank}")
trainer_type = params.get("trainer", {}).get("type", "default")
if trainer_type == "default":
# Special logic to instantiate backward-compatible trainer.
pieces = TrainerPieces.from_params(
params=params,
serialization_dir=serialization_dir,
recover=recover,
model=model,
embedding_sources_mapping=embedding_sources_mapping,
extend_vocab=extend_vocab,
)
trainer = Trainer.from_params(
model=pieces.model,
serialization_dir=serialization_dir,
iterator=pieces.iterator,
train_data=pieces.train_dataset,
validation_data=pieces.validation_dataset,
params=pieces.params,
validation_iterator=pieces.validation_iterator,
local_rank=process_rank,
)
evaluation_iterator = pieces.validation_iterator or pieces.iterator
evaluation_dataset = pieces.test_dataset
else:
if evaluate_on_test:
raise ValueError(
"--evaluate-on-test only works with the default Trainer. "
"If you're using the CallbackTrainer you can use a callback "
"to evaluate at Events.TRAINING_END; otherwise you'll have "
"to run allennlp evaluate separately.")
trainer = TrainerBase.from_params(params, serialization_dir, recover)
evaluation_dataset = None
evaluation_iterator = None
params.assert_empty("base train command")
try:
if distributed: # let the setup get ready for all the workers
dist.barrier()
metrics = trainer.train()
except KeyboardInterrupt:
# if we have completed an epoch, try to create a model archive.
if master and os.path.exists(
os.path.join(serialization_dir, _DEFAULT_WEIGHTS)):
logging.info(
"Training interrupted by the user. Attempting to create "
"a model archive using the current best epoch weights.")
archive_model(serialization_dir)
raise
if master:
if evaluation_dataset and evaluate_on_test:
logger.info(
"The model will be evaluated using the best epoch weights.")
test_metrics = evaluate(
trainer.model,
evaluation_dataset,
evaluation_iterator,
cuda_device=trainer.cuda_device,
batch_weight_key=batch_weight_key,
)
for key, value in test_metrics.items():
metrics["test_" + key] = value
elif evaluation_dataset:
logger.info(
"To evaluate on the test set after training, pass the "
"'evaluate_on_test' flag, or use the 'allennlp evaluate' command."
)
dump_metrics(os.path.join(serialization_dir, "metrics.json"),
metrics,
log=True)
if not distributed:
return trainer.model
return None # to make mypy happy
