def finetune(self, mixture_or_task_name, finetune_steps, pretrained_model_dir,
pretrained_checkpoint_step=-1, split="train"):
"""Finetunes a model from an existing checkpoint.
Args:
mixture_or_task_name: str, the name of the Mixture or Task to evaluate on.
Must be pre-registered in the global `TaskRegistry` or
`MixtureRegistry.`
finetune_steps: int, the number of additional steps to train for.
pretrained_model_dir: str, directory with pretrained model checkpoints and
operative config.
pretrained_checkpoint_step: int, checkpoint to initialize weights from. If
-1 (default), use the latest checkpoint from the pretrained model
directory.
split: str, the mixture/task split to finetune on.
"""
if pretrained_checkpoint_step == 0:
with gin.unlock_config():
gin.parse_config_file(_operative_config_path(pretrained_model_dir))
self.train(mixture_or_task_name, finetune_steps,
init_checkpoint=None,
split=split)
else:
if pretrained_checkpoint_step == -1:
checkpoint_step = utils.get_latest_checkpoint_from_dir(
pretrained_model_dir)
else:
checkpoint_step = pretrained_checkpoint_step
with gin.unlock_config():
gin.parse_config_file(_operative_config_path(pretrained_model_dir))
model_ckpt = "model.ckpt-" + str(checkpoint_step)
self.train(mixture_or_task_name, checkpoint_step + finetune_steps,
init_checkpoint=os.path.join(pretrained_model_dir, model_ckpt),
split=split)
def estimator(self, vocabulary, init_checkpoint=None, disable_tpu=False,
score_in_predict_mode=False):
if not self._tpu or disable_tpu:
with gin.unlock_config():
gin.bind_parameter("utils.get_variable_dtype.slice_dtype", "float32")
gin.bind_parameter(
"utils.get_variable_dtype.activation_dtype", "float32")
with gin.unlock_config():
gin.parse_config(self._gin_bindings)
return utils.get_estimator(
model_type=self._model_type,
vocabulary=vocabulary,
layout_rules=self._layout_rules,
mesh_shape=mtf.Shape([]) if disable_tpu else self._mesh_shape,
mesh_devices=None if disable_tpu else self._mesh_devices,
model_dir=self._model_dir,
batch_size=self.batch_size,
sequence_length=self._sequence_length,
autostack=self._autostack,
learning_rate_schedule=self._learning_rate_schedule,
keep_checkpoint_max=self._keep_checkpoint_max,
save_checkpoints_steps=self._save_checkpoints_steps,
optimizer=self._optimizer,
predict_fn=self._predict_fn,
variable_filter=self._variable_filter,
ensemble_inputs=self._ensemble_inputs,
use_tpu=None if disable_tpu else self._tpu,
tpu_job_name=self._tpu_job_name,
iterations_per_loop=self._iterations_per_loop,
cluster=self._cluster,
init_checkpoint=init_checkpoint,
score_in_predict_mode=score_in_predict_mode)
def load_model(instrument_model, audio_length):
# Build checkpoint path
# Assumes only one checkpoint in the folder, 'model.ckpt-[iter]`.
model_dir = os.path.join(CKPT_DIR,
"solo_%s_ckpt" % instrument_model.lower())
ckpt_files = [
f for f in tf.gfile.ListDirectory(model_dir) if "model.ckpt" in f
]
ckpt_name = ".".join(ckpt_files[0].split(".")[:2])
ckpt = os.path.join(model_dir, ckpt_name)
# Parse gin config
with gin.unlock_config():
gin_file = os.path.join(model_dir, "operative_config-0.gin")
gin.parse_config_file(gin_file, skip_unknown=True)
# Ensure dimensions sampling rates are equal
time_steps_train = gin.query_parameter("DefaultPreprocessor.time_steps")
n_samples_train = gin.query_parameter("Additive.n_samples")
hop_size = int(n_samples_train / time_steps_train)
time_steps = int(audio_length / hop_size)
n_samples = time_steps * hop_size
gin_params = [
"Additive.n_samples = {}".format(n_samples),
"FilteredNoise.n_samples = {}".format(n_samples),
"DefaultPreprocessor.time_steps = {}".format(time_steps),
]
with gin.unlock_config():
gin.parse_config(gin_params)
return ckpt, time_steps, n_samples
def load_model(self, model):
#model = 'Flute2' #@param ['Violin', 'Flute', 'Flute2', 'Trumpet', 'Tenor_Saxophone','Upload your own (checkpoint folder as .zip)']
MODEL = model
self.model_name = model
if model in ('Violin', 'Flute', 'Flute2', 'Trumpet',
'Tenor_Saxophone'):
# Pretrained models.
PRETRAINED_DIR = 'pretrained'
model_dir = PRETRAINED_DIR
gin_file = os.path.join(PRETRAINED_DIR, 'operative_config-0.gin')
# Parse gin config,
with gin.unlock_config():
gin.parse_config_file(gin_file, skip_unknown=True)
# Assumes only one checkpoint in the folder, 'ckpt-[iter]`.
ckpt_files = [f for f in tf.io.gfile.listdir(model_dir) if 'ckpt' in f]
ckpt_name = ckpt_files[0].split('.')[0]
ckpt = os.path.join(model_dir, ckpt_name)
# Ensure dimensions and sampling rates are equal
time_steps_train = gin.query_parameter(
'DefaultPreprocessor.time_steps')
n_samples_train = gin.query_parameter('Additive.n_samples')
hop_size = int(n_samples_train / time_steps_train)
print(self.audio.shape[1])
time_steps = int(self.audio.shape[1] / hop_size)
print(time_steps)
n_samples = time_steps * hop_size
print(n_samples)
gin_params = [
'Additive.n_samples = {}'.format(n_samples),
'FilteredNoise.n_samples = {}'.format(n_samples),
'DefaultPreprocessor.time_steps = {}'.format(time_steps),
]
with gin.unlock_config():
gin.parse_config(gin_params)
# Trim all input vectors to correct lengths
for key in ['f0_hz', 'f0_confidence', 'loudness_db']:
print(type(self.audio_features[key]))
self.audio_features[key] = self.audio_features[key][:time_steps]
print(self.audio_features['audio'].shape)
print(n_samples)
self.audio_features['audio'] = self.audio_features[
'audio'][:, :n_samples]
# Set up the model just to predict audio given new conditioning
self.model = ddsp.training.models.Autoencoder()
self.model.restore(ckpt)
# Build model by running a batch through it.
start_time = time.time()
_ = self.model(self.audio_features, training=False)
print('Restoring model took %.1f seconds' % (time.time() - start_time))
def _worker(self, root_dir, parameters, device_queue):
# sleep for random seconds to avoid crowded launching
try:
time.sleep(random.uniform(0, 3))
device = device_queue.get()
if self._conf.use_gpu:
os.environ["CUDA_VISIBLE_DEVICES"] = str(device)
else:
os.environ["CUDA_VISIBLE_DEVICES"] = "" # run on cpu
from alf.utils.common import set_per_process_memory_growth
set_per_process_memory_growth()
logging.set_verbosity(logging.INFO)
logging.info("parameters %s" % parameters)
with gin.unlock_config():
gin.parse_config(
['%s=%s' % (k, v) for k, v in parameters.items()])
train_eval(root_dir)
device_queue.put(device)
except Exception as e:
logging.info(e)
raise e
def latent1d(ctx, rows, cols, plot, filename, **kwargs):
"""Latent space traversal in 1D."""
add_gin(ctx, "config", ["evaluate/visual/latent1d.gin"])
parse(ctx, set_seed=True)
with gin.unlock_config():
gin.bind_parameter("disentangled.visualize.show.output.show_plot",
plot)
if filename is not None:
gin.bind_parameter("disentangled.visualize.show.output.filename",
filename)
if rows is not None:
gin.bind_parameter("disentangled.visualize.traversal1d.dimensions",
rows)
if cols is not None:
gin.bind_parameter("disentangled.visualize.traversal1d.steps",
cols)
dataset = ctx.obj["dataset"].pipeline()
disentangled.visualize.traversal1d(
ctx.obj["model"],
dataset,
dimensions=gin.REQUIRED,
offset=gin.REQUIRED,
skip_batches=gin.REQUIRED,
steps=gin.REQUIRED,
)
def score(self,
inputs,
targets,
scores_file=None,
checkpoint_steps=-1,
vocabulary=None):
"""Computes log-likelihood of target per example in targets.
Args:
inputs: optional - a string (filename), or a list of strings (inputs)
targets: a string (filename), or a list of strings (targets)
scores_file: str, path to write example scores to, one per line.
checkpoint_steps: int, list of ints, or None. If an int or list of ints,
inference will be run on the checkpoint files in `model_dir` whose
global steps are closest to the global steps provided. If None, run
inference continuously waiting for new checkpoints. If -1, get the
latest checkpoint from the model directory.
vocabulary: vocabularies.Vocabulary object to use for tokenization, or
None to use the default SentencePieceVocabulary.
"""
if checkpoint_steps == -1:
checkpoint_steps = _get_latest_checkpoint_from_dir(self._model_dir)
with gin.unlock_config():
gin.parse_config_file(_operative_config_path(self._model_dir))
# The following config setting ensures we do scoring instead of inference.
gin.bind_parameter("tpu_estimator_model_fn.score_in_predict_mode", True)
if vocabulary is None:
vocabulary = t5.data.get_default_vocabulary()
utils.score_from_strings(self.estimator(vocabulary), vocabulary,
self._model_type, self.batch_size,
self._sequence_length, self._model_dir,
checkpoint_steps, inputs, targets, scores_file)
def export(self, export_dir=None, checkpoint_step=-1, beam_size=1,
temperature=1.0,
sentencepiece_model_path=t5.data.DEFAULT_SPM_PATH):
"""Exports a TensorFlow SavedModel.
Args:
export_dir: str, a directory in which to export SavedModels. Will use
`model_dir` if unspecified.
checkpoint_step: int, checkpoint to export. If -1 (default), use the
latest checkpoint from the pretrained model directory.
beam_size: int, a number >= 1 specifying the number of beams to use for
beam search.
temperature: float, a value between 0 and 1 (must be 0 if beam_size > 1)
0.0 means argmax, 1.0 means sample according to predicted distribution.
sentencepiece_model_path: str, path to the SentencePiece model file to use
for decoding. Must match the one used during training.
"""
if checkpoint_step == -1:
checkpoint_step = _get_latest_checkpoint_from_dir(self._model_dir)
with gin.unlock_config():
gin.parse_config_file(_operative_config_path(self._model_dir))
gin.bind_parameter("Bitransformer.decode.beam_size", beam_size)
gin.bind_parameter("Bitransformer.decode.temperature", temperature)
gin.bind_parameter("utils.get_variable_dtype.slice_dtype", "float32")
gin.bind_parameter("utils.get_variable_dtype.activation_dtype", "float32")
vocabulary = t5.data.SentencePieceVocabulary(sentencepiece_model_path)
model_ckpt = "model.ckpt-" + str(checkpoint_step)
export_dir = export_dir or self._model_dir
utils.export_model(
self.estimator(vocabulary, disable_tpu=True), export_dir, vocabulary,
self._sequence_length, batch_size=self.batch_size,
checkpoint_path=os.path.join(self._model_dir, model_ckpt))
def configure_gin(self, ckpt):
"""Parse the model operative config with special streaming parameters."""
parse_operative_config(ckpt)
# Set streaming specific params.
preprocessor_ref = gin.query_parameter('Autoencoder.preprocessor')
preprocessor_str = preprocessor_ref.scoped_selector
time_steps = gin.query_parameter(f'{preprocessor_str}.time_steps')
n_samples = gin.query_parameter('Harmonic.n_samples')
if not isinstance(n_samples, int):
n_samples = gin.query_parameter('%n_samples')
samples_per_frame = int(n_samples / time_steps)
config = [
'Autoencoder.preprocessor = @F0PowerPreprocessor()',
'F0PowerPreprocessor.time_steps = 1',
f'Harmonic.n_samples = {samples_per_frame}',
f'FilteredNoise.n_samples = {samples_per_frame}',
]
# Remove reverb and crop processors.
processor_group_string = """ProcessorGroup.dag = [
(@synths.Harmonic(),
['amps', 'harmonic_distribution', 'f0_hz']),
(@synths.FilteredNoise(),
['noise_magnitudes']),
(@processors.Add(),
['filtered_noise/signal', 'harmonic/signal']),
]"""
config.append(processor_group_string)
with gin.unlock_config():
gin.parse_config(config)
def testSingleTrainingStepArchitectures(self,
use_predictor,
project_y=True,
self_supervision="none"):
parameters = {
"architecture": c.RESNET_BIGGAN_ARCH,
"lambda": 1,
"z_dim": 120,
}
with gin.unlock_config():
gin.bind_parameter("ModularGAN.conditional", True)
gin.bind_parameter("loss.fn", loss_lib.hinge)
gin.bind_parameter("S3GAN.use_predictor", use_predictor)
gin.bind_parameter("S3GAN.project_y", project_y)
gin.bind_parameter("S3GAN.self_supervision", self_supervision)
# Fake ImageNet dataset by overriding the properties.
dataset = datasets.get_dataset("imagenet_128")
model_dir = self._get_empty_model_dir()
run_config = tf.contrib.tpu.RunConfig(
model_dir=model_dir,
tpu_config=tf.contrib.tpu.TPUConfig(iterations_per_loop=1))
gan = S3GAN(dataset=dataset,
parameters=parameters,
model_dir=model_dir,
g_optimizer_fn=tf.train.AdamOptimizer,
g_lr=0.0002,
rotated_batch_fraction=2)
estimator = gan.as_estimator(run_config, batch_size=8, use_tpu=False)
estimator.train(gan.input_fn, steps=1)
def descartes_builder(name='out', params=[]):
print("Building grid search with parameters: ", params)
directory = os.path.join('grids', name)
if not os.path.exists(directory):
os.makedirs(directory)
all_values = []
all_params = []
for param in params:
values = gin.query_parameter(param)
all_params.append(param)
all_values.append(values)
descartes = itertools.product(*all_values)
i = 0
for one in descartes:
exp_directory = os.path.join(directory, str(i))
if not os.path.exists(exp_directory):
os.makedirs(exp_directory)
with gin.unlock_config():
for param_idx in range(len(all_params)):
gin.bind_parameter(all_params[param_idx], one[param_idx])
config_str = gin.config_str()
with open(os.path.join(exp_directory, 'config.gin'), 'w+') as f:
f.write(config_str)
i += 1
pass
def finetune(self, mixture_or_task_name, finetune_steps,
pretrained_model_dir, checkpoint_step):
"""Finetunes a model from an existing checkpoint.
Args:
mixture_or_task_name: str, the name of the Mixture or Task to evaluate on.
Must be pre-registered in the global `TaskRegistry` or
`MixtureRegistry.`
finetune_steps: int, the number of additional steps to train for.
pretrained_model_dir: str, directory with pretrained model checkpoints and
operative config.
checkpoint_step: int, checkpoint to initialize weights from. If -1, use
the latest checkpoint from the pretrained model directory.
"""
if checkpoint_step == -1:
checkpoint_step = get_latest_checkpoint_from_dir(
pretrained_model_dir)
model_ckpt = "model.ckpt-" + str(checkpoint_step)
with gin.unlock_config():
gin.parse_config_file(
os.path.join(pretrained_model_dir, "operative_config.gin"))
self.train(mixture_or_task_name,
checkpoint_step + finetune_steps,
init_checkpoint=os.path.join(pretrained_model_dir,
model_ckpt))
def configure_gin(self, ckpt):
"""Parse the model operative config with special streaming parameters."""
parse_operative_config(ckpt)
# Set streaming specific params.
time_steps = gin.query_parameter('F0PowerPreprocessor.time_steps')
n_samples = gin.query_parameter('Harmonic.n_samples')
samples_per_frame = int(n_samples / time_steps)
config = [
'F0PowerPreprocessor.time_steps = 1',
f'Harmonic.n_samples = {samples_per_frame}',
f'FilteredNoise.n_samples = {samples_per_frame}',
]
# Remove reverb processor.
processor_group_string = """ProcessorGroup.dag = [
(@synths.Harmonic(),
['amps', 'harmonic_distribution', 'f0_hz']),
(@synths.FilteredNoise(),
['noise_magnitudes']),
(@processors.Add(),
['filtered_noise/signal', 'harmonic/signal']),
]"""
config.append(processor_group_string)
with gin.unlock_config():
gin.parse_config(config)
def _runSingleTrainingStep(self, architecture, loss_fn, penalty_fn):
parameters = {
"architecture": architecture,
"lambda": 1,
"z_dim": 128,
}
with gin.unlock_config():
gin.bind_parameter("penalty.fn", penalty_fn)
gin.bind_parameter("loss.fn", loss_fn)
gin.bind_parameter("G.batch_norm_fn", evonorm_s0)
model_dir = self._get_empty_model_dir()
run_config = tf.contrib.tpu.RunConfig(
model_dir=model_dir,
tpu_config=tf.contrib.tpu.TPUConfig(iterations_per_loop=1))
dataset = datasets.get_dataset("cifar10")
gan = CLGAN(
dataset=dataset,
parameters=parameters,
model_dir=model_dir,
g_optimizer_fn=tf.train.AdamOptimizer,
g_lr=0.0002,
)
estimator = gan.as_estimator(run_config, batch_size=2, use_tpu=False)
estimator.train(gan.input_fn, steps=1)
def test_build_layer(self, kwarg_modules):
"""Tests if layer builds properly and produces outputs of correct shape."""
gin_config = (self.gin_config_kwarg_modules if kwarg_modules else
self.gin_config_dag_modules)
with gin.unlock_config():
gin.clear_config()
gin.parse_config(gin_config)
dag_layer = ConfigurableDAGLayer()
outputs = dag_layer(self.inputs)
self.assertIsInstance(outputs, dict)
z = outputs['bottleneck']['z_bottleneck']
x_rec = outputs['decoder']['reconstruction']
x_rec2 = outputs['out']['reconstruction']
# Confirm that layer generates correctly sized tensors.
self.assertEqual(outputs['test_data'].shape, self.x.shape)
self.assertEqual(outputs['inputs']['test_data'].shape, self.x.shape)
self.assertEqual(x_rec.shape, self.x.shape)
self.assertEqual(z.shape[-1], self.z_dims)
self.assertAllClose(x_rec, x_rec2)
# Confirm that variables are inherited by DAGLayer.
self.assertLen(dag_layer.trainable_variables, 6) # 3 weights, 3 biases.
def parse_gin(restore_dir):
"""Parse gin config from --gin_file, --gin_param, and the model directory."""
# Add user folders to the gin search path.
for gin_search_path in [GIN_PATH] + FLAGS.gin_search_path:
gin.add_config_file_search_path(gin_search_path)
# Parse gin configs, later calls override earlier ones.
with gin.unlock_config():
# Optimization defaults.
use_tpu = bool(FLAGS.tpu)
opt_default = 'base.gin' if not use_tpu else 'base_tpu.gin'
gin.parse_config_file(os.path.join('optimization', opt_default))
# Load operative_config if it exists (model has already trained).
operative_config = train_util.get_latest_operative_config(restore_dir)
if tf.io.gfile.exists(operative_config):
# Copy the config file from gstorage
helper_functions.copy_config_file_from_gstorage(
operative_config, LAST_OPERATIVE_CONFIG_PATH)
logging.info('Using operative config: %s', operative_config)
gin.parse_config_file(LAST_OPERATIVE_CONFIG_PATH,
skip_unknown=True)
# gin.parse_config_file(operative_config, skip_unknown=True)
# User gin config and user hyperparameters from flags.
gin.parse_config_files_and_bindings(FLAGS.gin_file,
FLAGS.gin_param,
skip_unknown=True)
def parse_gin(model_dir):
"""Parse gin config from --gin_file, --gin_param, and the model directory."""
# Add user folders to the gin search path.
for gin_search_path in [GIN_PATH] + FLAGS.gin_search_path:
gin.add_config_file_search_path(gin_search_path)
# Parse gin configs, later calls override earlier ones.
with gin.unlock_config():
# Optimization defaults.
use_tpu = bool(FLAGS.tpu)
opt_default = 'base.gin' if not use_tpu else 'base_tpu.gin'
gin.parse_config_file(os.path.join('optimization', opt_default))
# Load operative_config if it exists (model has already trained).
operative_config = os.path.join(model_dir, 'operative_config-0.gin')
if tf.io.gfile.exists(operative_config):
gin.parse_config_file(operative_config, skip_unknown=True)
# Only use the custom cumsum for TPUs.
gin.parse_config('ddsp.core.cumsum.use_tpu={}'.format(use_tpu))
# User gin config and user hyperparameters from flags.
gin.parse_config_files_and_bindings(FLAGS.gin_file,
FLAGS.gin_param,
skip_unknown=True)
def reset_config(self, new_config):
"""Resets trainer config for fast PBT implementation."""
with gin.unlock_config():
_parse_config_item(None, new_config['mutable_bindings'])
if new_config.get('reset_session'):
self.reset()
else:
if new_config.get('reset_optimizer'):
self._optimizer = None
self._model = None
if new_config.get('reset_problem'):
self._problem = None
self._generators = None
if new_config.get('reset_generators'):
self._generators = None
if new_config.get('reset_model'):
with tempfile.TemporaryDirectory() as tmp_dir:
checkpoint = self._save(tmp_dir)
self._model = None
self._restore(checkpoint)
self.config = new_config
self._reset_callbacks()
self._save_operative_config()
return True
def parse_gin_config(self, ckpt):
"""Parse the model operative config with special streaming parameters."""
with gin.unlock_config():
ckpt_dir = os.path.dirname(ckpt)
operative_config = train_util.get_latest_operative_config(ckpt_dir)
print(f'Parsing from operative_config {operative_config}')
gin.parse_config_file(operative_config, skip_unknown=True)
# Set streaming specific params.
# Remove reverb processor.
pg_string = """ProcessorGroup.dag = [
(@synths.Harmonic(),
['amps', 'harmonic_distribution', 'f0_hz']),
(@synths.FilteredNoise(),
['noise_magnitudes']),
(@processors.Add(),
['filtered_noise/signal', 'harmonic/signal']),
]"""
time_steps = gin.query_parameter('F0PowerPreprocessor.time_steps')
n_samples = gin.query_parameter('Harmonic.n_samples')
samples_per_frame = int(n_samples / time_steps)
gin.parse_config([
'F0PowerPreprocessor.time_steps=1',
f'Harmonic.n_samples={samples_per_frame}',
f'FilteredNoise.n_samples={samples_per_frame}',
pg_string,
])
def create_env(self):
"""Create the RL environment."""
"""Create an environment according to config."""
if 'env_config_file' in self.config:
with gin.unlock_config():
gin.parse_config_file(self.config['env_config_file'])
return load_env()
def _setup(self, config):
util.tf_init(gpus=None,
allow_growth=config.get('allow_growth', True),
eager=False)
logging.set_verbosity(config.get('verbosity', logging.INFO))
logging.info("calling setup")
config_files = config['config_files']
if isinstance(config_files, six.string_types):
config_files = [config_files]
with gin.unlock_config():
config_dir = config.get('config_dir')
if config_dir is None:
config_dir = util.get_config_dir()
util.parse_config(config_dir,
config_files,
config['bindings'],
finalize_config=False)
_parse_config_item(None, config['mutable_bindings'])
gin.finalize()
self._generators = None
self._problem = None
self._optimizer = None
self._model = None
self._reset_callbacks()
wp = config.get('initial_weights_path', None)
if wp is not None:
self.model.load_weights(wp)
self._save_operative_config()
def parse_gin_defaults_and_flags() -> None:
"""Parses all default gin files and those provided via flags."""
args = _parse_args() # TODO: should we require that args be passed in?
for path in [
*(args.gin_location_prefix or []),
Path(__file__).parent.parent.joinpath("config"),
pkg_resources.resource_filename("t5.models", "gin"),
pkg_resources.resource_filename("mesh_tensorflow.transformer",
"gin"),
]:
gin.add_config_file_search_path(path)
try:
# attempt to parse these first so they can be overridden later
gin.parse_config_file("defaults.gin")
gin.parse_config_file("operative_config.gin")
except IOError:
pass
gin.parse_config_files_and_bindings(args.gin_file, args.gin_param)
# make it so we don't have to specify a unique model_dir each time
model_dir = gin.query_parameter("utils.run.model_dir").format(
hostname=platform.node(),
timestamp=RUN_TIMESTAMP,
)
with gin.unlock_config():
gin.bind_parameter("utils.run.model_dir", model_dir)
tf_logging()
def evaluate_tune(log_dir,
split='validation',
verbose=True,
eval_batch_size=None):
from pointnet import train
log_dir = os.path.realpath(os.path.expanduser(os.path.expandvars(log_dir)))
paths = [os.path.join(log_dir, d) for d in tf.io.gfile.listdir(log_dir)]
op_configs = [d for d in paths if d.endswith('.gin')]
op_config = max(op_configs,
key=lambda x: parse_operative_config_path(x)[1])
dirs = [d for d in paths if tf.io.gfile.isdir(d)]
final_dir = max(dirs, key=lambda x: int(x.split('_')[-1]))
logging.info('Using operative_config {}, checkpoint {}'.format(
op_config, final_dir))
with gin.unlock_config():
gin.parse_config_file(op_config)
chkpt_callback = cb.ModelCheckpoint(final_dir)
if eval_batch_size is None:
eval_batch_size = blocks.batch_size()
train.evaluate(blocks.problem(),
blocks.model_fn(),
blocks.optimizer(),
eval_batch_size,
chkpt_callback,
split=split,
verbose=verbose)
def _worker(self, root_dir, parameters, device_queue):
# sleep for random seconds to avoid crowded launching
try:
time.sleep(random.uniform(0, 3))
device = device_queue.get()
if self._conf.use_gpu:
os.environ["CUDA_VISIBLE_DEVICES"] = str(device)
else:
os.environ["CUDA_VISIBLE_DEVICES"] = "" # run on cpu
if torch.cuda.is_available():
alf.set_default_device("cuda")
logging.set_verbosity(logging.INFO)
logging.info("Search parameters %s" % parameters)
with gin.unlock_config():
gin.parse_config(
['%s=%s' % (k, v) for k, v in parameters.items()])
gin.parse_config(
"TrainerConfig.confirm_checkpoint_upon_crash=False")
train_eval(FLAGS.ml_type, root_dir)
device_queue.put(device)
except Exception as e:
logging.info(traceback.format_exc())
raise e
def fixed(ctx, batch_size, filename, rows, cols, plot, verbose, **kwargs):
"""View/save images of dataset given a fixed latent factor."""
dataset = ctx.obj['dataset']
add_gin(ctx, 'config', ['evaluate/dataset/{}.gin'.format(dataset)])
parse(ctx)
with gin.unlock_config():
gin.bind_parameter('disentangled.visualize.show.output.show_plot',
plot)
gin.bind_parameter('disentangled.visualize.show.output.filename',
filename)
if rows is not None:
gin.bind_parameter('disentangled.visualize.fixed_factor_data.rows',
rows)
if cols is not None:
gin.bind_parameter('disentangled.visualize.fixed_factor_data.cols',
cols)
num_values_per_factor = disentangled.dataset.get(
dataset).num_values_per_factor
dataset = disentangled.dataset.get(dataset).supervised()
fixed, _ = disentangled.metric.utils.fixed_factor_dataset(
dataset, batch_size, num_values_per_factor)
disentangled.visualize.fixed_factor_data(fixed,
rows=gin.REQUIRED,
cols=gin.REQUIRED,
verbose=verbose)
def parse_gin(restore_dir):
"""Parse gin config from --gin_file, --gin_param, and the model directory."""
# Add user folders to the gin search path.
for gin_search_path in [GIN_PATH] + FLAGS.gin_search_path:
gin.add_config_file_search_path(gin_search_path)
# Parse gin configs, later calls override earlier ones.
with gin.unlock_config():
# Optimization defaults.
use_tpu = bool(FLAGS.tpu)
opt_default = 'base.gin' if not use_tpu else 'base_tpu.gin'
gin.parse_config_file(os.path.join('optimization', opt_default))
eval_default = 'eval/basic.gin'
gin.parse_config_file(eval_default)
# Load operative_config if it exists (model has already trained).
operative_config = train_util.get_latest_operative_config(restore_dir)
if tf.io.gfile.exists(operative_config):
logging.info('Using operative config: %s', operative_config)
operative_config = cloud.make_file_paths_local(
operative_config, GIN_PATH)
gin.parse_config_file(operative_config, skip_unknown=True)
# User gin config and user hyperparameters from flags.
gin_file = cloud.make_file_paths_local(FLAGS.gin_file, GIN_PATH)
gin.parse_config_files_and_bindings(gin_file,
FLAGS.gin_param,
skip_unknown=True)
def eval(self, mixture_or_task_name, checkpoint_steps=None, summary_dir=None,
split="validation"):
"""Evaluate the model on the given Mixture or Task.
Args:
mixture_or_task_name: str, the name of the Mixture or Task to evaluate on.
Must be pre-registered in the global `TaskRegistry` or
`MixtureRegistry.`
checkpoint_steps: int, list of ints, or None. If an int or list of ints,
evaluation will be run on the checkpoint files in `model_dir` whose
global steps are closest to the global steps provided. If None, run eval
continuously waiting for new checkpoints. If -1, get the latest
checkpoint from the model directory.
summary_dir: str, path to write TensorBoard events file summaries for
eval. If None, use model_dir/eval_{split}.
split: str, the mixture/task split to evaluate on.
"""
if checkpoint_steps == -1:
checkpoint_steps = _get_latest_checkpoint_from_dir(self._model_dir)
vocabulary = t5.models.mesh_transformer.get_vocabulary(mixture_or_task_name)
dataset_fn = functools.partial(
t5.models.mesh_transformer.mesh_eval_dataset_fn,
mixture_or_task_name=mixture_or_task_name,
)
with gin.unlock_config():
gin.parse_config_file(_operative_config_path(self._model_dir))
utils.eval_model(self.estimator(vocabulary), vocabulary,
self._sequence_length, self.batch_size, split,
self._model_dir, dataset_fn, summary_dir, checkpoint_steps)
def visual(ctx, rows, cols, plot, filename, **kwargs):
"""Qualitative evaluation of output."""
parse(ctx, set_seed=True)
with gin.unlock_config():
gin.bind_parameter("disentangled.visualize.show.output.show_plot",
plot)
if filename is not None:
gin.bind_parameter("disentangled.visualize.show.output.filename",
filename)
if rows is not None:
gin.bind_parameter("disentangled.visualize.reconstructed.rows",
rows)
if cols is not None:
gin.bind_parameter("disentangled.visualize.reconstructed.cols",
cols)
dataset = ctx.obj["dataset"].pipeline()
disentangled.visualize.reconstructed(ctx.obj["model"],
dataset,
rows=gin.REQUIRED,
cols=gin.REQUIRED)
def export(self, export_dir=None, checkpoint_step=-1, beam_size=1,
temperature=1.0, vocabulary=None):
"""Exports a TensorFlow SavedModel.
Args:
export_dir: str, a directory in which to export SavedModels. Will use
`model_dir` if unspecified.
checkpoint_step: int, checkpoint to export. If -1 (default), use the
latest checkpoint from the pretrained model directory.
beam_size: int, a number >= 1 specifying the number of beams to use for
beam search.
temperature: float, a value between 0 and 1 (must be 0 if beam_size > 1)
0.0 means argmax, 1.0 means sample according to predicted distribution.
vocabulary: vocabularies.Vocabulary object to use for tokenization, or
None to use the default SentencePieceVocabulary.
Returns:
The string path to the exported directory.
"""
if checkpoint_step == -1:
checkpoint_step = _get_latest_checkpoint_from_dir(self._model_dir)
with gin.unlock_config():
gin.parse_config_file(_operative_config_path(self._model_dir))
gin.bind_parameter("Bitransformer.decode.beam_size", beam_size)
gin.bind_parameter("Bitransformer.decode.temperature", temperature)
if vocabulary is None:
vocabulary = t5.data.get_default_vocabulary()
model_ckpt = "model.ckpt-" + str(checkpoint_step)
export_dir = export_dir or self._model_dir
return utils.export_model(
self.estimator(vocabulary, disable_tpu=True), export_dir, vocabulary,
self._sequence_length, batch_size=self.batch_size,
checkpoint_path=os.path.join(self._model_dir, model_ckpt))
def load_operative_gin_configurations(operative_config_dir):
"""Load operative Gin configurations from the given directory."""
gin_log_file = operative_config_path(operative_config_dir)
with gin.unlock_config():
gin.parse_config_file(gin_log_file)
gin.finalize()
logging.info('Operative Gin configurations loaded from %s.', gin_log_file)
