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))