def test_pipeline_split(self):
def initial_model_1():
model_input = keras.Input(shape=(32, 32, 3))
model_output = keras.layers.MaxPooling2D(
name='test_pipeline_split_layer1')(model_input)
model_output_1 = keras.layers.Conv2D(
filters=32, kernel_size=3,
name='test_pipeline_split_layer2')(model_output)
model_output_2 = keras.layers.Conv2D(
filters=32, kernel_size=3,
name='test_pipeline_split_layer3')(model_output)
model_output = keras.layers.Add(name='test_pipeline_split_layer4')(
[model_output_1, model_output_2])
model_output = keras.layers.Flatten(
name='test_pipeline_split_layer5')(model_output)
return keras.Model(model_input, model_output)
def expected_model_1():
model_input = keras.Input(shape=(32, 32, 3))
with ipu.keras.PipelineStage(0):
model_output = keras.layers.MaxPooling2D()(model_input)
model_output_1 = keras.layers.Conv2D(
filters=32, kernel_size=3)(model_output)
with ipu.keras.PipelineStage(1):
model_output_2 = keras.layers.Conv2D(
filters=32, kernel_size=3)(model_output)
model_output = keras.layers.Add()(
[model_output_1, model_output_2])
with ipu.keras.PipelineStage(2):
model_output = keras.layers.Flatten()(model_output)
return keras.Model(model_input, model_output)
train_strategy = ipu.ipu_strategy.IPUStrategy()
with train_strategy.scope():
model = initial_model_1()
pipelined_model = ModelFactory.configure_model(
model=model,
gradient_accumulation_count=1,
pipeline_splits=[
'test_pipeline_split_layer3', 'test_pipeline_split_layer5'
],
device_mapping=[],
pipeline_schedule='Grouped',
available_memory_proportion=[])
expected_assignments = expected_model_1(
).get_pipeline_stage_assignment()
pipelined_assignments = pipelined_model.get_pipeline_stage_assignment(
)
for expected_assignment, pipelined_assignment in zip(
expected_assignments, pipelined_assignments):
assert (expected_assignment.layer.__class__.name ==
pipelined_assignment.layer.__class__.name)
assert (expected_assignment.pipeline_stage ==
pipelined_assignment.pipeline_stage)
def test_invalid_id_of_device_mapping(self):
train_strategy = ipu.ipu_strategy.IPUStrategy()
with train_strategy.scope():
model = ModelFactory.create_model(model_name='toy_model',
weights=None,
input_shape=(28, 28, 1),
classes=10)
with self.assertRaises(DimensionError):
model = ModelFactory.configure_model(
model=model,
gradient_accumulation_count=1,
pipeline_splits=['conv2d_1', 'flatten'],
device_mapping=[1, 2, 3],
pipeline_schedule='Grouped',
available_memory_proportion=[])
with train_strategy.scope():
# Create an instance of the model
model = ModelFactory.create_model(
model_name=model_name,
input_shape=img_shape,
classes=num_classes,
accelerator_side_preprocessing_fn=
accelerator_side_preprocess_train_fn,
eight_bit_transfer=eight_bit_transfer)
model = ModelFactory.configure_model(
model=model,
gradient_accumulation_count=batch_config.
gradient_accumulation_count,
pipeline_splits=pipeline_splits,
device_mapping=device_mapping,
pipeline_schedule=pipeline_schedule,
available_memory_proportion=available_memory_proportion,
optimizer_state_offloading=optimizer_state_offloading)
if training:
# prepare the learning rate scheduler
lr_outfeed_queue = ipu.ipu_outfeed_queue.IPUOutfeedQueue(
outfeed_mode=ipu.ipu_outfeed_queue.IPUOutfeedMode.LAST)
lr_scheduler = get_lr_scheduler(
scheduler_name=lr_schedule,
schedule_params=lr_schedule_params,
warmup_params=lr_warmup_params,
global_batch_size=batch_config.global_batch_size,
weight_updates_per_epoch=weight_updates_per_epoch,