def test_training():
"""Verify only that no error raised."""
config = SmartDict()
config.NETWORK_CLASS = LMBiSeNet
config.DATASET_CLASS = DummyCamvid
config.IS_DEBUG = False
config.IMAGE_SIZE = [128, 160]
config.BATCH_SIZE = 2
config.TEST_STEPS = 1
config.MAX_STEPS = 2
config.SAVE_CHECKPOINT_STEPS = 1
config.KEEP_CHECKPOINT_MAX = 5
config.SUMMARISE_STEPS = 1
config.IS_PRETRAIN = False
config.TASK = Tasks.SEMANTIC_SEGMENTATION
# network model config
config.NETWORK = SmartDict()
config.NETWORK.OPTIMIZER_CLASS = tf.train.AdamOptimizer
config.NETWORK.OPTIMIZER_KWARGS = {"learning_rate": 0.001}
config.NETWORK.IMAGE_SIZE = config.IMAGE_SIZE
config.NETWORK.BATCH_SIZE = config.BATCH_SIZE
config.NETWORK.DATA_FORMAT = "NHWC"
# daasegt config
config.DATASET = SmartDict()
config.DATASET.PRE_PROCESSOR = Resize(config.IMAGE_SIZE)
config.DATASET.BATCH_SIZE = config.BATCH_SIZE
config.DATASET.DATA_FORMAT = "NHWC"
environment.init("test_lm_bisenet")
prepare_dirs(recreate=True)
start_training(config, profile_step=1)
def test_training():
"""Test only that no error raised."""
config = EasyDict()
config.NETWORK_CLASS = YoloV1
config.DATASET_CLASS = Pascalvoc2007
config.IS_DEBUG = False
config.IMAGE_SIZE = [70, 70]
config.BATCH_SIZE = 4
config.TEST_STEPS = 1
config.MAX_STEPS = 2
config.SAVE_CHECKPOINT_STEPS = 1
config.KEEP_CHECKPOINT_MAX = 5
config.SUMMARISE_STEPS = 1
config.IS_PRETRAIN = False
config.TASK = Tasks.OBJECT_DETECTION
# network model config
config.NETWORK = EasyDict()
config.NETWORK.IMAGE_SIZE = config.IMAGE_SIZE
config.NETWORK.BATCH_SIZE = config.BATCH_SIZE
# daasegt config
config.DATASET = EasyDict()
config.DATASET.PRE_PROCESSOR = ResizeWithGtBoxes(config.IMAGE_SIZE)
config.DATASET.BATCH_SIZE = config.BATCH_SIZE
environment.init("test_yolov_1")
prepare_dirs(recreate=True)
start_training(config)
def test_save_json():
experiment_id = "test_save_json"
environment.init(experiment_id)
setup_test_environment()
test_dict = {
"model_name": "save_json",
"image_size_height": 128,
"image_size_width": 64,
"num_classes": 3,
"parameters": "test_node",
"flops": "test_flops",
}
if not os.path.exists(environment.EXPERIMENT_DIR):
os.makedirs(environment.EXPERIMENT_DIR)
_save_json(name=test_dict["model_name"],
image_size=(test_dict["image_size_height"],
test_dict["image_size_width"]),
num_classes=test_dict["num_classes"],
node_param_dict=test_dict["parameters"],
node_flops_dict=test_dict["flops"])
output_file = os.path.join(
environment.EXPERIMENT_DIR,
"{}_profile.json".format(test_dict["model_name"]))
with open(output_file, 'r') as fp:
file_data = json.load(fp)
assert os.path.isfile(output_file)
assert test_dict == file_data
def run(input_dir, output_dir, experiment_id, config_file, restore_path,
save_images):
environment.init(experiment_id)
config = config_util.load_from_experiment()
if config_file:
config = config_util.merge(config, config_util.load(config_file))
if not os.path.isdir(input_dir):
raise FileNotFoundError(
"Input directory not found: '{}'".format(input_dir))
if restore_path is None:
restore_file = search_restore_filename(environment.CHECKPOINTS_DIR)
restore_path = os.path.join(environment.CHECKPOINTS_DIR, restore_file)
print("Restore from {}".format(restore_path))
if not os.path.exists("{}.index".format(restore_path)):
raise FileNotFoundError(
"Checkpoint file not found: '{}'".format(restore_path))
print("---- start predict ----")
_run(input_dir, output_dir, config, restore_path, save_images)
print("---- end predict ----")
def test_training():
"""Test only that no error raised."""
config = EasyDict()
config.NETWORK_CLASS = YoloV2
config.DATASET_CLASS = Pascalvoc2007
config.IS_DEBUG = False
config.IMAGE_SIZE = [128, 160]
config.BATCH_SIZE = 2
config.TEST_STEPS = 1
config.MAX_STEPS = 2
config.SAVE_CHECKPOINT_STEPS = 1
config.KEEP_CHECKPOINT_MAX = 5
config.SUMMARISE_STEPS = 1
config.IS_PRETRAIN = False
config.TASK = Tasks.OBJECT_DETECTION
# network model config
config.NETWORK = EasyDict()
config.NETWORK.OPTIMIZER_CLASS = tf.train.AdamOptimizer
config.NETWORK.OPTIMIZER_KWARGS = {"learning_rate": 0.001}
config.NETWORK.IMAGE_SIZE = config.IMAGE_SIZE
config.NETWORK.BATCH_SIZE = config.BATCH_SIZE
config.NETWORK.DATA_FORMAT = "NHWC"
# dataset config
config.DATASET = EasyDict()
config.DATASET.PRE_PROCESSOR = ResizeWithGtBoxes(config.IMAGE_SIZE)
config.DATASET.BATCH_SIZE = config.BATCH_SIZE
config.DATASET.DATA_FORMAT = "NHWC"
environment.init("test_yolo_v2")
prepare_dirs(recreate=True)
start_training(config, profile_step=1)
def test_training():
"""Test only no error raised."""
config = EasyDict()
config.NETWORK_CLASS = Darknet
config.DATASET_CLASS = Dummy
config.IS_DEBUG = False
config.IMAGE_SIZE = [28, 14]
config.BATCH_SIZE = 2
config.TEST_STEPS = 1
config.MAX_STEPS = 2
config.SAVE_CHECKPOINT_STEPS = 1
config.KEEP_CHECKPOINT_MAX = 5
config.SUMMARISE_STEPS = 1
config.IS_PRETRAIN = False
config.TASK = Tasks.CLASSIFICATION
# network model config
config.NETWORK = EasyDict()
config.NETWORK.OPTIMIZER_CLASS = tf.train.AdamOptimizer
config.NETWORK.OPTIMIZER_KWARGS = {"learning_rate": 0.001}
config.NETWORK.IMAGE_SIZE = config.IMAGE_SIZE
config.NETWORK.BATCH_SIZE = config.BATCH_SIZE
# daasegt config
config.DATASET = EasyDict()
config.DATASET.PRE_PROCESSOR = Resize(config.IMAGE_SIZE)
config.DATASET.BATCH_SIZE = config.BATCH_SIZE
environment.init("test_darknet")
prepare_dirs(recreate=True)
start_training(config)
def run(experiment_id,
restore_path=None,
image_size=(None, None),
image=DEFAULT_INFERENCE_TEST_DATA_IMAGE,
config_file=None):
environment.init(experiment_id)
config = config_util.load_from_experiment()
if config_file:
config = config_util.merge(config, config_util.load(config_file))
config.BATCH_SIZE = 1
config.NETWORK.BATCH_SIZE = 1
config.DATASET.BATCH_SIZE = 1
if list(image_size) != [None, None]:
config.IMAGE_SIZE = list(image_size)
config.NETWORK.IMAGE_SIZE = list(image_size)
# override pre processes image size.
if config.PRE_PROCESSOR:
config.PRE_PROCESSOR.set_image_size(image_size)
# override post processes image size.
if config.POST_PROCESSOR:
config.POST_PROCESSOR.set_image_size(image_size)
print("Override IMAGE_SIZE", config.IMAGE_SIZE)
executor.init_logging(config)
config_util.display(config)
return _export(config, restore_path, image)
def run(experiment_id, restore_path, config_file, bit, unquant_layers):
if config_file is None and experiment_id is None:
raise Exception("config_file or experiment_id are required")
if experiment_id:
environment.init(experiment_id)
config = config_util.load_from_experiment()
if config_file:
config = config_util.merge(config, config_util.load(config_file))
if restore_path is None:
restore_file = executor.search_restore_filename(environment.CHECKPOINTS_DIR)
restore_path = os.path.join(environment.CHECKPOINTS_DIR, restore_file)
if not os.path.exists("{}.index".format(restore_path)):
raise Exception("restore file {} dont exists.".format(restore_path))
else:
experiment_id = "profile"
environment.init(experiment_id)
config = config_util.load(config_file)
config.BATCH_SIZE = 1
config.NETWORK.BATCH_SIZE = 1
config.DATASET.BATCH_SIZE = 1
executor.init_logging(config)
config_util.display(config)
_profile(config, restore_path, bit, unquant_layers)
def test_example_object_detection_config_yaml():
"""Test that export config and meta yaml from example object_detection config python."""
config_file = os.path.join("..", "blueoil", "configs", "example",
"object_detection.py")
config = _load_py(config_file)
config_yaml = os.path.join("..", "blueoil", "configs", "example",
"object_detection.yaml")
config_meta = os.path.join("..", "blueoil", "configs", "example",
"object_detection_meta.yaml")
environment.init("test_example_object_detection_config_yaml")
saved_config, saved_meta = save_yaml(environment.EXPERIMENT_DIR, config)
with open(config_yaml) as f:
expected = f.read()
with open(saved_config) as f:
data = f.read()
assert expected == data
with open(config_meta) as f:
expected = f.read()
with open(saved_meta) as f:
data = f.read()
assert expected == data
def main(output_file_base, metrics_keys, steps, experiment_id):
environment.init(experiment_id)
output(
environment.TENSORBOARD_DIR,
environment.EXPERIMENT_DIR,
metrics_keys,
steps,
output_file_base="metrics",
)
def main(config_file, experiment_id, restore_path, output_dir):
environment.init(experiment_id)
config = config_util.load_from_experiment()
if config_file:
config = config_util.merge(config, config_util.load(config_file))
executor.init_logging(config)
config_util.display(config)
evaluate(config, restore_path, output_dir)
def test_init_with_experiment_dir(self):
experiment_id = "experiment001"
experiment_dir = os.path.join(environment.OUTPUT_DIR, experiment_id)
environment.init(experiment_dir)
assert environment.EXPERIMENT_ID == experiment_id
assert environment.EXPERIMENT_DIR == experiment_dir
assert environment.TENSORBOARD_DIR == os.path.join(
experiment_dir,
"tensorboard",
)
assert environment.CHECKPOINTS_DIR == os.path.join(
experiment_dir,
"checkpoints",
)
def run(config_file, experiment_id, recreate):
environment.init(experiment_id)
config = config_util.load(config_file)
if horovod_util.is_enabled():
horovod_util.setup()
if horovod_util.is_rank0():
config_util.display(config)
executor.init_logging(config)
executor.prepare_dirs(recreate)
config_util.copy_to_experiment_dir(config_file)
config_util.save_yaml(environment.EXPERIMENT_DIR, config)
start_training(config)
def test_training():
"""Test only no error raised."""
config = EasyDict()
config.NETWORK_CLASS = LmSinglePoseV1Quantize
config.DATASET_CLASS = MscocoSinglePersonKeypoints
config.IS_DEBUG = False
config.IMAGE_SIZE = [160, 160]
config.BATCH_SIZE = 2
config.TEST_STEPS = 1
config.MAX_STEPS = 2
config.SAVE_CHECKPOINT_STEPS = 1
config.KEEP_CHECKPOINT_MAX = 5
config.SUMMARISE_STEPS = 1
config.IS_PRETRAIN = False
config.IS_DISTRIBUTION = False
config.TASK = Tasks.KEYPOINT_DETECTION
# network model config
config.NETWORK = EasyDict()
config.NETWORK.OPTIMIZER_CLASS = tf.train.AdamOptimizer
config.NETWORK.OPTIMIZER_KWARGS = {"learning_rate": 0.001}
config.NETWORK.IMAGE_SIZE = config.IMAGE_SIZE
config.NETWORK.BATCH_SIZE = config.BATCH_SIZE
config.NETWORK.ACTIVATION_QUANTIZER = linear_mid_tread_half_quantizer
config.NETWORK.ACTIVATION_QUANTIZER_KWARGS = {
'bit': 2,
'max_value': 2.0
}
config.NETWORK.WEIGHT_QUANTIZER = binary_channel_wise_mean_scaling_quantizer
config.NETWORK.WEIGHT_QUANTIZER_KWARGS = {}
# daasegt config
config.DATASET = EasyDict()
config.DATASET.PRE_PROCESSOR = Sequence([
ResizeWithJoints(image_size=config.IMAGE_SIZE),
JointsToGaussianHeatmap(image_size=config.IMAGE_SIZE, stride=2),
DivideBy255()])
config.DATASET.BATCH_SIZE = config.BATCH_SIZE
environment.init("test_lm_single_pose_v1")
prepare_dirs(recreate=True)
start_training(config, profile_step=1)
def main(network, dataset, config_file, experiment_id, restore_path,
output_dir):
environment.init(experiment_id)
config = config_util.load_from_experiment()
if config_file:
config = config_util.merge(config, config_util.load(config_file))
if network:
network_class = module_loader.load_network_class(network)
config.NETWORK_CLASS = network_class
if dataset:
dataset_class = module_loader.load_dataset_class(dataset)
config.DATASET_CLASS = dataset_class
executor.init_logging(config)
config_util.display(config)
evaluate(config, restore_path, output_dir)
def main(model):
if model == "yolov2":
weight_file = 'inputs/yolo-voc.weights'
experiment_id = "convert_weight_from_darknet/yolo_v2"
config_file = "configs/convert_weight_from_darknet/yolo_v2.py"
if model == "darknet19":
weight_file = 'inputs/darknet19_448.weights'
experiment_id = "convert_weight_from_darknet/darknet19"
config_file = "configs/convert_weight_from_darknet/darknet19.py"
recreate = True
environment.init(experiment_id)
executor.prepare_dirs(recreate)
config = config_util.load(config_file)
config_util.display(config)
config_util.copy_to_experiment_dir(config_file)
convert(config, weight_file)
def test_training():
"""Test only that no error raised."""
config = EasyDict()
config.NETWORK_CLASS = LmResnetQuantize
config.DATASET_CLASS = Dummy
config.IS_DEBUG = False
config.IMAGE_SIZE = [32, 32]
config.BATCH_SIZE = 2
config.TEST_STEPS = 1
config.MAX_STEPS = 2
config.SAVE_CHECKPOINT_STEPS = 1
config.KEEP_CHECKPOINT_MAX = 5
config.SUMMARISE_STEPS = 1
config.IS_PRETRAIN = False
config.TASK = Tasks.CLASSIFICATION
# network model config
config.NETWORK = EasyDict()
config.NETWORK.OPTIMIZER_CLASS = tf.train.AdamOptimizer
config.NETWORK.OPTIMIZER_KWARGS = {"learning_rate": 0.001}
config.NETWORK.IMAGE_SIZE = config.IMAGE_SIZE
config.NETWORK.BATCH_SIZE = config.BATCH_SIZE
config.NETWORK.ACTIVATION_QUANTIZER = linear_mid_tread_half_quantizer
config.NETWORK.ACTIVATION_QUANTIZER_KWARGS = {
'bit': 2,
'max_value': 2
}
config.NETWORK.WEIGHT_QUANTIZER = binary_mean_scaling_quantizer
config.NETWORK.WEIGHT_QUANTIZER_KWARGS = {}
# dataset config
config.DATASET = EasyDict()
config.DATASET.PRE_PROCESSOR = Resize(config.IMAGE_SIZE)
config.DATASET.BATCH_SIZE = config.BATCH_SIZE
environment.init("test_lm_resnet_quantize")
prepare_dirs(recreate=True)
start_training(config)
def test_set_output_dir_after_init(self):
experiment_id = "experiment001"
environment.init(experiment_id)
path = "./output"
environment.set_output_dir(path)
abs_path = os.path.abspath(path)
assert environment.EXPERIMENT_DIR == os.path.join(
abs_path,
experiment_id,
)
assert environment.TENSORBOARD_DIR == os.path.join(
abs_path,
experiment_id,
"tensorboard",
)
assert environment.CHECKPOINTS_DIR == os.path.join(
abs_path,
experiment_id,
"checkpoints",
)
def test_training():
"""Test only no error raised."""
config = SmartDict()
config.NETWORK_CLASS = YoloV2Quantize
config.DATASET_CLASS = Pascalvoc2007
config.IS_DEBUG = False
config.IMAGE_SIZE = [128, 160]
config.BATCH_SIZE = 2
config.TEST_STEPS = 1
config.MAX_STEPS = 2
config.SAVE_CHECKPOINT_STEPS = 1
config.KEEP_CHECKPOINT_MAX = 5
config.SUMMARISE_STEPS = 1
config.IS_PRETRAIN = False
config.TASK = Tasks.OBJECT_DETECTION
# network model config
config.NETWORK = SmartDict()
config.NETWORK.OPTIMIZER_CLASS = tf.train.AdamOptimizer
config.NETWORK.OPTIMIZER_KWARGS = {"learning_rate": 0.001}
config.NETWORK.IMAGE_SIZE = config.IMAGE_SIZE
config.NETWORK.BATCH_SIZE = config.BATCH_SIZE
config.NETWORK.ACTIVATION_QUANTIZER = linear_mid_tread_half_quantizer
config.NETWORK.ACTIVATION_QUANTIZER_KWARGS = {'bit': 2, 'max_value': 2.0}
config.NETWORK.WEIGHT_QUANTIZER = binary_channel_wise_mean_scaling_quantizer
config.NETWORK.WEIGHT_QUANTIZER_KWARGS = {}
# daasegt config
config.DATASET = SmartDict()
config.DATASET.PRE_PROCESSOR = ResizeWithGtBoxes(config.IMAGE_SIZE)
config.DATASET.BATCH_SIZE = config.BATCH_SIZE
environment.init("test_yolov_2_quantize")
prepare_dirs(recreate=True)
start_training(config, profile_step=1)
def run(network, dataset, config_file, experiment_id, recreate):
environment.init(experiment_id)
config = config_util.load(config_file)
if network:
network_class = module_loader.load_network_class(network)
config.NETWORK_CLASS = network_class
if dataset:
dataset_class = module_loader.load_dataset_class(dataset)
config.DATASET_CLASS = dataset_class
if horovod_util.is_enabled():
horovod_util.setup()
if horovod_util.is_rank0():
config_util.display(config)
executor.init_logging(config)
executor.prepare_dirs(recreate)
config_util.copy_to_experiment_dir(config_file)
config_util.save_yaml(environment.EXPERIMENT_DIR, config)
start_training(config)
def _run(config_file, experiment_id, restore_path, image_size, step_size, cpu):
if experiment_id:
environment.init(experiment_id)
config = config_util.load_from_experiment()
if config_file:
config = config_util.merge(config, config_util.load(config_file))
if restore_path is None:
restore_file = executor.search_restore_filename(
environment.CHECKPOINTS_DIR)
restore_path = os.path.join(environment.CHECKPOINTS_DIR,
restore_file)
if not os.path.exists("{}.index".format(restore_path)):
raise Exception(
"restore file {} dont exists.".format(restore_path))
else:
experiment_id = "measure_latency"
environment.init(experiment_id)
config = config_util.load(config_file)
config.BATCH_SIZE = 1
config.NETWORK.BATCH_SIZE = 1
config.DATASET.BATCH_SIZE = 1
if list(image_size) != [None, None]:
config.IMAGE_SIZE = list(image_size)
config.NETWORK.IMAGE_SIZE = list(image_size)
# override pre processes image size.
if config.PRE_PROCESSOR:
config.PRE_PROCESSOR.set_image_size(image_size)
# override post processes image size.
if config.POST_PROCESSOR:
config.POST_PROCESSOR.set_image_size(image_size)
print("Override IMAGE_SIZE", config.IMAGE_SIZE)
executor.init_logging(config)
config_util.display(config)
overall_times, only_network_times = _measure_time(config, restore_path,
step_size)
overall_times = np.array(overall_times)
only_network_times = np.array(only_network_times)
# list of physical_device_desc
devices = [
device.physical_device_desc
for device in device_lib.list_local_devices()
if device.physical_device_desc
]
message = """
---- measure latency result ----
total number of execution (number of samples): {}
network: {}
use gpu by network: {}
image size: {}
devices: {}
* overall (include pre-post-process which execute on cpu)
total time: {:.4f} msec
latency
mean (SD=standard deviation): {:.4f} (SD={:.4f}) msec, min: {:.4f} msec, max: {:.4f} msec
FPS
mean (SD=standard deviation): {:.4f} (SD={:.4f}), min: {:.4f}, max: {:.4f}
* network only (exclude pre-post-process):
total time: {:.4f} msec
latency
mean (SD=standard deviation): {:.4f} (SD={:.4f}) msec, min: {:.4f} msec, max: {:.4f} msec
FPS
mean (SD=standard deviation): {:.4f} (SD={:.4f}), min: {:.4f}, max: {:.4f}
---- measure latency result ----
""".format(
step_size,
config.NETWORK_CLASS.__name__,
not cpu,
config.IMAGE_SIZE,
devices,
# overall
np.sum(overall_times) * 1000,
# latency
np.mean(overall_times) * 1000,
np.std(overall_times) * 1000,
np.min(overall_times) * 1000,
np.max(overall_times) * 1000,
# FPS
np.mean(1 / overall_times),
np.std(1 / overall_times),
np.min(1 / overall_times),
np.max(1 / overall_times),
# network only
np.sum(only_network_times) * 1000,
# latency
np.mean(only_network_times) * 1000,
np.std(only_network_times) * 1000,
np.min(only_network_times) * 1000,
np.max(only_network_times) * 1000,
# FPS
np.mean(1 / only_network_times),
np.std(1 / only_network_times),
np.min(1 / only_network_times),
np.max(1 / only_network_times),
)
print(message)
def test_profile_train_step():
"""Asserts that the saved profile matches expectation for a simple example."""
environment.init("test_executor")
prepare_dirs(recreate=True)
with tf.compat.v1.Session() as sess:
w = tf.Variable([[5.0, 3.0, 2.9, -4.0, 0.0]])
v = tf.Variable([[0.21, -2.70, 0.94, 3.82, -3.65],
[5.0, 3.0, 2.9, -4.0, 0.0],
[1.96, -2.2, 0.42, -1.26, -1.06],
[-1.55, 4.56, -4.71, -2.43, 4.55],
[-3.11, 3.78, -3.45, 2.18, -4.45]])
z = tf.matmul(
w, v) # z is [[27.933998, -29.119999, 33.458, 13.166, -39.524002]]
sess.run(tf.compat.v1.global_variables_initializer())
step = 0
options = tf.compat.v1.RunOptions(
trace_level=tf.compat.v1.RunOptions.FULL_TRACE)
run_meta = tf.compat.v1.RunMetadata()
sess.run(z, options=options, run_metadata=run_meta)
profile_train_step(step, sess, run_meta)
expected_memory = textwrap.dedent("""\
Doc:
scope: The nodes in the model graph are organized by their names, which is hierarchical like filesystem.
requested bytes: The memory requested by the operation, accumulatively.
Profile:
node name | requested bytes
_TFProfRoot
MatMul
Variable
Variable_1
""").splitlines()
expected_memory.sort()
expected_timeline = [{
"args": {
"name": "MatMul",
"op": "MatMul"
},
"cat": "Op",
"dur": 267,
"name": "MatMul",
"ph": "X",
"pid": 1,
"tid": 0,
"ts": 0
}, {
"args": {
"name": "Variable",
"op": "Variable"
},
"cat": "Op",
"dur": 4,
"name": "Variable",
"ph": "X",
"pid": 1,
"tid": 0,
"ts": 267
}, {
"args": {
"name": "Variable_1",
"op": "Variable_1"
},
"cat": "Op",
"dur": 20,
"name": "Variable_1",
"ph": "X",
"pid": 1,
"tid": 0,
"ts": 271
}, {
"args": {
"name": "_TFProfRoot",
"op": "_TFProfRoot"
},
"cat": "Op",
"dur": 291,
"name": "_TFProfRoot",
"ph": "X",
"pid": 0,
"tid": 0,
"ts": 0
}, {
"args": {
"name": "Scope:0"
},
"name": "process_name",
"ph": "M",
"pid": 0
}, {
"args": {
"name": "Scope:1"
},
"name": "process_name",
"ph": "M",
"pid": 1
}]
train_memory_path = os.path.join(environment.EXPERIMENT_DIR,
"training_profile_memory")
with open(train_memory_path) as train_memory_file:
saved_data = train_memory_file.read().splitlines()
saved_data.sort()
for idx, line in enumerate(saved_data):
assert line.startswith(expected_memory[idx])
train_timeline_path = os.path.join(environment.EXPERIMENT_DIR,
"training_profile_timeline_step")
with open("{}_{}".format(train_timeline_path,
step)) as train_timeline_file:
saved_data = json.load(train_timeline_file)["traceEvents"]
saved_data.sort(key=lambda op: op["name"])
for op1, op2 in zip(expected_timeline, saved_data):
assert op1["args"] == op2["args"]
# Generally, timeline values are different each run, so just check the keys match.
assert op1.keys() == op2.keys()
