def __enter__(self):
if self.custom_ops:
nnef._register_custom_ops(self.key, self.custom_ops)
if self.custom_shapes:
nnef._register_custom_shapes(self.custom_shapes)
if self.deferred_shapes:
nnef._register_deferred_shapes(self.deferred_shapes)
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if self.custom_ops:
nnef._unregister_custom_ops(self.key)
if self.custom_shapes:
custom_shapes = self.custom_shapes.copy()
for k in custom_shapes.keys():
custom_shapes[k] = None
nnef._register_custom_shapes(custom_shapes)
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import nnef
def shuffle_shape(op, args, shapes):
shapes[args['output']] = shapes[args['input']]
nnef._register_custom_ops(
"shuffle",
"fragment shuffle<?>( input: tensor<?>, groups: integer ) -> ( output: tensor<?> );"
)
nnef._register_custom_shapes({"shuffle": shuffle_shape})
graph = nnef.parse_string("""
version 1.0;
graph Net( input ) -> ( output )
{
input = external(shape = [1,3,224,224]);
filter = variable(shape = [32,3,5,5], label = 'conv/filter');
conv = conv(input, filter);
output = shuffle(conv, groups = 4);
}
""")
print(
nnef.format_graph(graph.name, graph.inputs, graph.outputs,
graph.operations))