Commit c2a1ecd
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -292,21 +292,31 @@ def _Net_batch(self, blobs): | |||
| 292 | 292 | padding]) | |
| 293 | 293 | yield padded_batch | |
| 294 | 294 | ||
| 295 | - | ||
| 296 | - class _Net_IdNameWrapper: | ||
| 297 | - """ | ||
| 298 | - A simple wrapper that allows the ids propery to be accessed as a dict | ||
| 299 | - indexed by names. Used for top and bottom names | ||
| 295 | + def _Net_get_id_name(func, field): | ||
| 300 | 296 | """ | |
| 301 | - def __init__(self, net, func): | ||
| 302 | - self.net, self.func = net, func | ||
| 297 | + Generic property that maps func to the layer names into an OrderedDict. | ||
| 298 | + | ||
| 299 | + Used for top_names and bottom_names. | ||
| 303 | 300 | ||
| 304 | - def __getitem__(self, name): | ||
| 305 | - # Map the layer name to id | ||
| 306 | - ids = self.func(self.net, list(self.net._layer_names).index(name)) | ||
| 307 | - # Map the blob id to name | ||
| 308 | - id_to_name = list(self.net.blobs) | ||
| 309 | - return [id_to_name[i] for i in ids] | ||
| 301 | + Parameters | ||
| 302 | + ---------- | ||
| 303 | + func: function id -> [id] | ||
| 304 | + field: implementation field name (cache) | ||
| 305 | + | ||
| 306 | + Returns | ||
| 307 | + ------ | ||
| 308 | + A one-parameter function that can be set as a property. | ||
| 309 | + """ | ||
| 310 | + @property | ||
| 311 | + def get_id_name(self): | ||
| 312 | + if not hasattr(self, field): | ||
| 313 | + id_to_name = list(self.blobs) | ||
| 314 | + res = OrderedDict([(self._layer_names[i], | ||
| 315 | + [id_to_name[j] for j in func(self, i)]) | ||
| 316 | + for i in range(len(self.layers))]) | ||
| 317 | + setattr(self, field, res) | ||
| 318 | + return getattr(self, field) | ||
| 319 | + return get_id_name | ||
| 310 | 320 | ||
| 311 | 321 | # Attach methods to Net. | |
| 312 | 322 | Net.blobs = _Net_blobs | |
@@ -320,5 +330,5 @@ def __getitem__(self, name): | |||
| 320 | 330 | Net._batch = _Net_batch | |
| 321 | 331 | Net.inputs = _Net_inputs | |
| 322 | 332 | Net.outputs = _Net_outputs | |
| 323 | - Net.top_names = property(lambda n: _Net_IdNameWrapper(n, Net._top_ids)) | ||
| 324 | - Net.bottom_names = property(lambda n: _Net_IdNameWrapper(n, Net._bottom_ids)) | ||
| 333 | + Net.top_names = _Net_get_id_name(Net._top_ids, "_top_names") | ||
| 334 | + Net.bottom_names = _Net_get_id_name(Net._bottom_ids, "_bottom_names") | ||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -3,6 +3,7 @@ | |||
| 3 | 3 | import os | |
| 4 | 4 | import numpy as np | |
| 5 | 5 | import six | |
| 6 | + from collections import OrderedDict | ||
| 6 | 7 | ||
| 7 | 8 | import caffe | |
| 8 | 9 | ||
@@ -78,6 +79,18 @@ def test_inputs_outputs(self): | |||
| 78 | 79 | self.assertEqual(self.net.inputs, []) | |
| 79 | 80 | self.assertEqual(self.net.outputs, ['loss']) | |
| 80 | 81 | ||
| 82 | + def test_top_bottom_names(self): | ||
| 83 | + self.assertEqual(self.net.top_names, | ||
| 84 | + OrderedDict([('data', ['data', 'label']), | ||
| 85 | + ('conv', ['conv']), | ||
| 86 | + ('ip', ['ip']), | ||
| 87 | + ('loss', ['loss'])])) | ||
| 88 | + self.assertEqual(self.net.bottom_names, | ||
| 89 | + OrderedDict([('data', []), | ||
| 90 | + ('conv', ['data']), | ||
| 91 | + ('ip', ['conv']), | ||
| 92 | + ('loss', ['ip', 'label'])])) | ||
| 93 | + | ||
| 81 | 94 | def test_save_and_read(self): | |
| 82 | 95 | f = tempfile.NamedTemporaryFile(mode='w+', delete=False) | |
| 83 | 96 | f.close() | |
0 commit comments