utils.py

importre
fromurllib.parseimporturlparse
fromshutilimportrmtree
importlogging
importos
frompathlibimportPath
importsys
importtensorflow.compat.v1astf
importtensorflow.compat.v2astf2
importmesh_tensorflowasmtf
importmesh_tensorflow.auto_mtf
fromdata.encodersimportfetch_encoder
importre

defsetup_logging(args):
Path("logs").mkdir(exist_ok=True)
tf.logging.set_verbosity(logging.INFO)
tf.get_logger().propagate=False# Remove double log on console
name=os.path.splitext(os.path.basename(args.model))[0]
handlers= [
logging.FileHandler(f"logs/{name}.log"),
logging.StreamHandler(sys.stdout)
 ]
logger=logging.getLogger("tensorflow")
logger.handlers=handlers
returnlogger


defget_batch_size(params):
returnparams[f"{params['mode']}_batch_size"]


defadd_mode_to_params(params, mode):
ifmode==tf.estimator.ModeKeys.PREDICT:
params["mode"] ="predict"
elifmode==tf.estimator.ModeKeys.EVAL:
params["mode"] ="eval"
elifmode==tf.estimator.ModeKeys.TRAIN:
params["mode"] ="train"
else:
raiseValueError(f"Invalid mode {mode}")
returnparams


defsimd_mesh_setup(params, mesh_shape, layout_rules):
"""Constructs SimdMesh function - instructions on how to evenly split tensors across all TPU cores"""

num_hosts=params["context"].num_hosts
host_placement_fn=params["context"].tpu_host_placement_function
device_list= [host_placement_fn(host_id=i) foriinrange(num_hosts)]
tf.logging.info(f"device_list = {device_list}")

# TODO: Better estimation of replica cache size?
replica_cache_size=300*1000000# 300M per replica

# Worker 0 caches all the TPU binaries
worker0_mem=replica_cache_size*params["context"].num_replicas
devices_memory_usage= [worker0_mem] + [0] * (num_hosts-1)
var_placer=mtf.utils.BalancedVariablePlacer(device_list, devices_memory_usage)
mesh_devices= [""] *mesh_shape.size
mesh_impl=mtf.simd_mesh_impl.SimdMeshImpl(
mesh_shape, layout_rules, mesh_devices, params["context"].device_assignment)

returnvar_placer, mesh_impl


defremove_batch_from_layout(layout):
"""
 The tf-mesh layout splits across batch size, remove it.
 Useful for prediction steps, when you no longer want large batches.

 :param layout: string describing tf-mesh layout
 :return: layout minus batch dimension
 """
layout=layout.split(',')
ret_layout=""
foriinlayout:
if"batch"ini:
pass
else:
ret_layout+=f"{i},"
returnret_layout[:-1]


defyes_or_no(question):
whileTrue:
reply=str(input(question+' (y/n): ')).lower().strip()
ifreply[:1] =='y':
returnTrue
ifreply[:1] =='n':
returnFalse


defremove_gs_or_filepath(path):
parsed_url=urlparse(path)
ifparsed_url.scheme=="gs":
os.system(f"gsutil rm -rf {path}")
return
rmtree(path)


defsave_config(params_dict, logdir):
print(f"Saving config to {logdir}")
text="{\n\n"
total_params=len(params_dict)
forcount, keyinenumerate(params_dict):
config_value=str(params_dict[key])
ifre.search('[a-zA-Z]', config_value):
ifconfig_value.lower() !='true':
ifconfig_value.lower() !='false':
ifconfig_value[0] !='[':
# TODO: Making a manual exception for parsing epsilon right now since it's the only number in
# scientific notation. Should fix this.
ifkey!="epsilon":
config_value=f'"{config_value}"'
ifcount==total_params-1:
text+=f'"{str(key)}"'+' : '+config_value+'\n\n'
else:
text+=f'"{str(key)}"'+' : '+config_value+',\n\n'
text+='\n\n}'
sess=tf.InteractiveSession()
summary_op=tf.summary.text("run_config", tf.convert_to_tensor(text))
summary_writer=tf.summary.FileWriter(f"{logdir}/config", sess.graph)
text=sess.run(summary_op)
summary_writer.add_summary(text, 0)
summary_writer.flush()
summary_writer.close()
tf.reset_default_graph()
print('Done!')


defexpand_attention_types_params(params_list):
newlist= []
foriteminparams_list:
for_inrange(item[1]):
newlist.extend(item[0])
returnnewlist


defget_n_trainable_vars(graph):
"""
 Gets number of trainable vars in a MTF model.

 :param graph: Mesh-Tensorflow graph
 :return: None
 """
total_parameters=0
forvariableingraph.trainable_variables:
shape=variable.shape.dims
variable_parameters=1
fordiminshape:
variable_parameters*=dim.size
total_parameters+=variable_parameters
print(f"\n\nN TRAINABLE VARS:\n{total_parameters:,}\n\n")


defprint_dim_names(graph):
"""
 Print names of all Dimensions
 :param graph: Mesh-Tensorflow graph
 :return: None
 """
all_dim_names= []
forvariableingraph.all_variables:
names=variable.shape.dimension_names
all_dim_names.append(names)

# Print all dim names in graph & write to file
all_dim_names= [itemforsublistinall_dim_namesforiteminsublist] # Flatten all dims
unique_dims=list(set(all_dim_names))
print("ALL DIM NAMES:")
fordim_nameinunique_dims:
print(dim_name)
print('\n')


defget_graph_info(graph):
"""
 Wrapper fn that calculates number of trainable vars in an MTF graph & prints all dim_names to file
 TODO: how to get un-trainable dim-names too, batch etc.

 :param graph: Mesh-Tensorflow graph
 :return: None
 """
get_n_trainable_vars(graph)
print_dim_names(graph)


defloss_denominator(targets, num_microbatches):
"""Denominator applied to losses.

 This is usually the size of the targets tensor (omitting ensemble
 dimensions). Alternatively, it is an override value passed to the
 class constructor.

 Args:
 targets: a mtf.Tensor
 num_microbatches: an integer - greater than one if the step has been
 serialized into multiple microbatches to save memory.
 Returns:
 a float
 """
ret=float(targets.shape.size) *num_microbatches
returnfloat(ret)

defcheck_dataset(input_fn, params, global_step=None):
tf.enable_eager_execution()
ifglobal_stepisnotNone:
dataset=input_fn(params, global_step=global_step)
else:
dataset=input_fn(params)
dataset_iter=dataset.make_one_shot_iterator()
tensor, _=next(dataset_iter)
enc=fetch_encoder(params)

forpintensor[:1]:
txt=enc.decode(p)

print('-'*50)
print(txt[:500], '\n\n...\n\n', txt[-500:])
print('-'*50)
exit()

defauto_layout(graph, mesh_shape, logits, loss):
layout_rules=mtf.auto_mtf.layout(graph, mesh_shape, [logits, loss])
print(f"Auto-selected layout:\n{layout_rules}\nRe-initialize graph with selected layout")
quit() 

defauto_layout_and_mesh_shape(graph, num_cores, logits, loss):
layout_rules, mesh_shape=mtf.auto_mtf.layout_and_mesh_shape(graph, num_cores,
 [logits, loss], max_mesh_shape_dimensions=4)
print(f"Num cores:\n{num_cores}\nAuto-selected layout:\n{layout_rules}\nAuto-selected mesh shape:\n{mesh_shape}" \
f"\nRe-initialize graph with selected layout & mesh shape")
quit() 

defcreate_host_call(model_dir):
"""Construct a host_call writing scalar summaries.

 Borrowed from t2t.

 Args:
 model_dir: String containing path to train
 Returns:
 (fn, args) Pair to be called by TPUEstimator as the host_call.
 """

graph=tf.get_default_graph()
# A list of (name, lowered tensor) tuples
summaries=graph.get_collection(mtf.utils.SCALAR_SUMMARIES_COLLECTION_KEY)

defmaybe_cast(tensor):
asserttensor.shape.is_compatible_with([]), tensor.name
iftensor.dtype==tf.int64:
returntf.to_int32(tensor)
iftensor.dtype==tf.bfloat16:
returntf.cast(tensor, tf.float32)
returntensor

reshaped_tensors= [tf.reshape(maybe_cast(t), [1]) for_, tinsummaries]

# When no supported summaries are found, don't create host_call. Otherwise,
# TPU outfeed queue would enqueue global_step while host_call doesn't dequeue
# it, eventually causing hang.
ifnotreshaped_tensors:
returnNone

defhost_call_fn(global_step, *args):
"""Training host call. Creates scalar summaries for training metrics."""
# This function is executed on the CPU and should not directly reference
# any Tensors in the rest of the `model_fn`. To pass Tensors from the
# model to the `model_fn`, provide as part of the `host_call`.
global_step=tf.cast(global_step[0], tf.int64)
withtf2.summary.create_file_writer(model_dir).as_default():
# We cannot directly use any tensor from summaries, because each
# tensor here must be a concat of multiple tensors from all shards.
# Therefore, we rely on the assumption that args wil have the same
# length as summaries, and all tensors in args will have the same
# order of self._tup_summaries.
assertlen(args) ==len(summaries)
fori, tensorinenumerate(args):
name=summaries[i][0]
tf2.summary.scalar(name, tf.reduce_mean(tensor), step=global_step)
returntf.summary.all_v2_summary_ops()

global_step_t=tf.reshape(tf.to_int32(tf.train.get_global_step()), [1])
returnhost_call_fn, [global_step_t] +reshaped_tensors


defnatural_sort(l): 
convert=lambdatext: int(text) iftext.isdigit() elsetext.lower() 
alphanum_key=lambdakey: [ convert(c) forcinre.split('([0-9]+)', key) ] 
returnsorted(l, key=alphanum_key)