def test_prune_dangling(self):
"""Tests dangling vertices are pruned."""
model1 = model_spec.ModelSpec(
np.array([[0, 1, 1, 0],
[0, 0, 0, 0],
[0, 0, 0, 1],
[0, 0, 0, 0]]),
[1, 2, 3, 4])
assert model1.valid_spec
assert np.array_equal(model1.matrix,
np.array([[0, 1, 0],
[0, 0, 1],
[0, 0, 0]]))
assert model1.ops == [1, 3, 4]
model2 = model_spec.ModelSpec(
np.array([[0, 0, 1, 0],
[0, 0, 0, 1],
[0, 0, 0, 1],
[0, 0, 0, 0]]),
[1, 2, 3, 4])
assert model2.valid_spec
assert np.array_equal(model2.matrix,
np.array([[0, 1, 0],
[0, 0, 1],
[0, 0, 0]]))
assert model2.ops == [1, 3, 4]
def test_prune_islands(self):
"""Tests isolated components are pruned."""
model1 = model_spec.ModelSpec(
np.array([[0, 1, 0, 0],
[0, 0, 0, 1],
[0, 0, 0, 0],
[0, 0, 0, 0]]),
[1, 2, 3, 4])
assert model1.valid_spec
assert np.array_equal(model1.matrix,
np.array([[0, 1, 0],
[0, 0, 1],
[0, 0, 0]]))
assert model1.ops == [1, 2, 4]
model2 = model_spec.ModelSpec(
np.array([[0, 1, 0, 0, 0],
[0, 0, 0, 0, 1],
[0, 0, 0, 1, 0],
[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0]]),
[1, 2, 3, 4, 5])
assert model2.valid_spec
assert np.array_equal(model2.matrix,
np.array([[0, 1, 0],
[0, 0, 1],
[0, 0, 0]]))
assert model2.ops == [1, 2, 5]
def test_prune_noop(self):
"""Tests graphs which require no pruning."""
model1 = model_spec.ModelSpec(
np.array([[0, 1, 0],
[0, 0, 1],
[0, 0, 0]]),
[0, 0, 0])
assert model1.valid_spec
assert np.array_equal(model1.original_matrix, model1.matrix)
assert model1.original_ops == model1.original_ops
model2 = model_spec.ModelSpec(
np.array([[0, 1, 1],
[0, 0, 1],
[0, 0, 0]]),
[0, 0, 0])
assert model2.valid_spec
assert np.array_equal(model2.original_matrix, model2.matrix)
assert model2.original_ops == model2.ops
model3 = model_spec.ModelSpec(
np.array([[0, 1, 1, 0],
[0, 0, 0, 1],
[0, 0, 0, 1],
[0, 0, 0, 0]]),
[0, 0, 0, 0])
assert model3.valid_spec
assert np.array_equal(model3.original_matrix, model3.matrix)
assert model3.original_ops == model3.ops
def create_inception_resnet_spec(config):
"""Construct an Inception-ResNet like spec.
This spec is very similar to the InceptionV2 module with an added
residual connection except that there is an extra projection in front of the
max pool. The overall network filter counts and module counts do not match
the actual source model.
Args:
config: config dict created by config.py.
Returns:
ModelSpec object.
"""
spec = model_spec.ModelSpec(
np.array([[0, 1, 1, 1, 0, 1, 1], [0, 0, 0, 0, 0, 0, 1],
[0, 0, 0, 0, 0, 0, 1], [0, 0, 0, 0, 1, 0, 0],
[0, 0, 0, 0, 0, 0, 1], [0, 0, 0, 0, 0, 0, 1],
[0, 0, 0, 0, 0, 0, 0]]),
[
'input', 'conv1x1-bn-relu', 'conv3x3-bn-relu', 'conv3x3-bn-relu',
'conv3x3-bn-relu', 'maxpool3x3', 'output'
])
config['num_stacks'] = 3
config['num_modules_per_stack'] = 3
config['stem_filter_size'] = 128
return spec
def test_prune_disconnected(self):
"""Tests graphs where with no input to output path are marked invalid."""
model1 = model_spec.ModelSpec(
np.array([[0, 0],
[0, 0]]),
[0, 0])
assert not model1.valid_spec
model2 = model_spec.ModelSpec(
np.array([[0, 1, 0, 0],
[0, 0, 0, 0],
[0, 0, 0, 1],
[0, 0, 0, 0]]),
[1, 2, 3, 4])
assert not model2.valid_spec
model3 = model_spec.ModelSpec(
np.array([[0, 0, 0, 0],
[0, 0, 1, 0],
[0, 0, 0, 0],
[0, 0, 0, 0]]),
[1, 2, 3, 4])
assert not model3.valid_spec
def _evaluate_work_unit(self, index):
"""Runs the evaluation of the model at the specified index.
The index records the current index of the work unit being evaluated. Each
worker will only compute the work units with index modulo total_workers
equal to the worker_id.
Args:
index: int index into total work units.
"""
if self.remainders:
assert self.ordered_keys is None
model_id = self.remainders[index][0]
model_repeat = self.remainders[index][1]
else:
model_id = self.ordered_keys[index % self.num_models]
model_repeat = index // self.num_models + 1
matrix, labels = self.models[model_id]
matrix = np.array(matrix)
# Re-label to config['available_ops']
labels = (['input'] +
[self.config['available_ops'][lab]
for lab in labels[1:-1]] + ['output'])
spec = model_spec.ModelSpec(matrix, labels)
assert spec.valid_spec
assert np.sum(spec.matrix) <= self.config['max_edges']
# Split the directory into 16^2 roughly equal subdirectories
model_dir = os.path.join(self.output_dir, model_id[:2], model_id,
'repeat_%d' % model_repeat)
try:
meta = evaluate.train_and_evaluate(spec, self.config, model_dir)
except evaluate.AbortError:
# After hitting the retry limit, the job will continue to the next work
# unit. These failed jobs may need to be re-run at a later point.
return
# Write data to model_dir
output_file = os.path.join(model_dir, RESULTS_FILE)
with tf.gfile.Open(output_file, 'w') as f:
json.dump(meta, f, cls=NumpyEncoder)
# Delete some files to reclaim space
self._clean_model_dir(model_dir)
def create_resnet50_spec(config):
"""Construct a ResNet-50-like spec.
The main difference is that there is an extra projection layer before the
conv1x1 whereas the original ResNet doesn't have this. This increases the
parameter count of this version slightly.
Args:
config: config dict created by config.py.
Returns:
ModelSpec object.
"""
spec = model_spec.ModelSpec(np.array([[0, 1, 1], [0, 0, 1], [0, 0, 0]]),
['input', 'bottleneck3x3', 'output'])
config['num_stacks'] = 3
config['num_modules_per_stack'] = 6
config['stem_filter_size'] = 128
return spec
def create_best_nasbench_spec(config):
"""Construct the best spec in the NASBench dataset w.r.t. mean test accuracy.
Args:
config: config dict created by config.py.
Returns:
ModelSpec object.
"""
spec = model_spec.ModelSpec(
np.array([[0, 1, 1, 0, 0, 1, 1], [0, 0, 0, 0, 0, 1, 0],
[0, 0, 0, 1, 0, 0, 0], [0, 0, 0, 0, 1, 0, 0],
[0, 0, 0, 0, 0, 1, 0], [0, 0, 0, 0, 0, 0, 1],
[0, 0, 0, 0, 0, 0, 0]]),
[
'input', 'conv1x1-bn-relu', 'conv3x3-bn-relu', 'maxpool3x3',
'conv3x3-bn-relu', 'conv3x3-bn-relu', 'output'
])
config['num_stacks'] = 3
config['num_modules_per_stack'] = 3
config['stem_filter_size'] = 128
return spec