in visualize/generate.py [0:0]
def extract_iteration(epoch_data):
data_by_iter = []
current_iter = []
broadcast_count = 0
for item in epoch_data:
if item.comm_type == CommType.epoch_end:
if current_iter:
data_by_iter.append(current_iter)
break
if broadcast_count < 2:
if item.comm_type == CommType.broadcast:
broadcast_count += 1
current_iter.append(item)
else:
if item.comm_type == CommType.broadcast:
data_by_iter.append(current_iter)
current_iter = [item]
broadcast_count = 1
else:
current_iter.append(item)
return data_by_iter