def to_microbatched(model, micro_batch_size, num_micro_batches, num_batches, num_test_batches):
rank = tnt.get_rank()
partition_generator = pgen.GraphPartitionGenerator(model)
rank_mapper = rmapper.RankMapper(num_ranks = tnt.get_size(),
pipeline_graph = partition_generator.get_pipeline_graph())
partition_id = rank_mapper.get_partition_for_rank(rank)
partition_graph = partition_generator.get_partition_graph(partition_id)
partition_info = pinfo.PartitionInfo(partition_id = partition_id,
partition_graph = partition_graph)
core_model_builder = cm_builder.CoreModelBuilder(model, partition_id, partition_graph)
core_model = core_model_builder.get_model()
connection_table = rank_mapper.get_connections_for_rank(rank)
pipeline_communicator = tnt.PipelineCommunicator(connection_table, num_micro_batches)
shared_model_builder = shared.SharedModelBuilder(partition_info, core_model,
pipeline_communicator, micro_batch_size)
shared_model = shared_model_builder.get_model()
microbatched_model_builder = microbatched.MicrobatchedModelBuilder(partition_info, shared_model,
micro_batch_size, num_micro_batches)
ds = load_microbatched_datasets(micro_batch_size, num_micro_batches,
num_batches, num_test_batches, partition_info)
pipeline_communicator.setup_infrastructure(micro_batch_size)
return microbatched_model_builder, ds
def test_partition_core_models(self, model_and_partitions):
num_micro_batches = 1
model, partition_gen, expected_num_partitions, _, expected_model_gen = model_and_partitions
rank_mapper = rmapper.RankMapper(
num_ranks=expected_num_partitions,
pipeline_graph=partition_gen.get_pipeline_graph())
for rank in range(expected_num_partitions):
partition_id = rank_mapper.get_partition_for_rank(rank)
partition_graph = partition_gen.get_partition_graph(partition_id)
cm_builder = core_model_builder.CoreModelBuilder(
model, partition_id, partition_graph)
core_model = cm_builder.get_model()
reference_core_model = expected_model_gen(rank)
utils.check_model_configuration_identical(core_model,
reference_core_model)
utils.compare_weights(core_model.get_weights(),
reference_core_model.get_weights(), 1e-6)
def __init__(self, model, group, partition_generator, rank_mapper,
num_pipeline_stages = None):
super().__init__(model = model, group = group)
self._model_name = model.name
self.built = False
self.compile_properties = None
self.num_pipeline_stages = num_pipeline_stages
connection_table = rank_mapper.get_connections_for_rank(self.rank)
self.pipeline_communicator = tnt.PipelineCommunicator(connection_table, self.num_pipeline_stages)
self.initialized = False
partition_id = rank_mapper.get_partition_for_rank(self.rank)
partition_graph = partition_generator.get_partition_graph(partition_id)
self.partition_info = pinfo.PartitionInfo(partition_id, partition_graph)
core_model_builder = cm_builder.CoreModelBuilder(model, partition_id,
partition_graph)
self.model = core_model_builder.get_model()
self.nano_batch_size = None
self.built = False
def test_core_model_outputs(self, model_and_partitions):
num_micro_batches = 1
model, partition_gen, expected_num_partitions, _, expected_model_gen = model_and_partitions
rank_mapper = rmapper.RankMapper(
num_ranks=expected_num_partitions,
pipeline_graph=partition_gen.get_pipeline_graph())
for rank in range(expected_num_partitions):
tf.keras.backend.clear_session()
partition_id = rank_mapper.get_partition_for_rank(rank)
partition_graph = partition_gen.get_partition_graph(partition_id)
cm_builder = core_model_builder.CoreModelBuilder(
model, partition_id, partition_graph)
core_model = cm_builder.get_model()
core_output_names = [i.name for i in core_model.outputs]
tf.keras.backend.clear_session()
reference_core_model = expected_model_gen(rank)
reference_output_names = [
i.name for i in reference_core_model.outputs
]
assert core_output_names == reference_output_names
strides=(1, 1),
padding='valid',
activation='relu')(x)
x = keras.layers.MaxPooling2D(pool_size=(3, 3),
strides=(2, 2),
padding='valid')(y)
y = keras.layers.MaxPooling2D(pool_size=(3, 3),
strides=(2, 2),
padding='valid')(y)
x = keras.layers.Concatenate()([x, y])
x = keras.layers.Flatten()(x)
x = keras.layers.Dense(512, activation='relu')(x)
outputs = keras.layers.Dense(10, activation='softmax')(x)
model = keras.Model(inputs=inputs, outputs=outputs)
return model
model = alexnet_model_generator()
partition_generator = pgen.GraphPartitionGenerator(model)
nranks = 3
rank_mapper = rank_mapper.RankMapper(partition_generator.get_pipeline_graph(),
nranks)
for rank in range(nranks):
pprint.pprint(f"Model for rank {rank}")
core_model_builder = core_model.CoreModelBuilder(partition_generator,
rank_mapper, rank)
model = core_model_builder.get_model()
model.summary()