| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, 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 mxnet as mx |
| |
| @mx.util.use_np |
| def test_save(): |
| class MyBlock(mx.gluon.Block): |
| def __init__(self, **kwargs): |
| super(MyBlock, self).__init__(**kwargs) |
| self.layers = [] |
| def add(self, block): |
| self.layers.append(block) |
| self.register_child(block) |
| def forward(self, x, *args): |
| out = (x,) + args |
| for block in self._children.values(): |
| out = block()(*out) |
| return out |
| |
| def createNet(): |
| inside = MyBlock() |
| dense = mx.gluon.nn.Dense(10) |
| inside.add(dense) |
| net = MyBlock() |
| net.add(inside) |
| net.add(mx.gluon.nn.Dense(10)) |
| return net |
| |
| # create and initialize model |
| net1 = createNet() |
| net1.initialize() |
| # hybridize (the hybridizeable blocks, ie. the Dense layers) |
| net1.hybridize() |
| x = mx.np.zeros((1,10)) |
| out1 = net1(x) |
| |
| # save hybridized model |
| net1.save('MyModel') |
| |
| # create a new model, uninitialized |
| net2 = createNet() |
| # reload hybridized model |
| net2.load('MyModel') |
| net2.hybridize() |
| # run inference again |
| out2 = net2(x) |
| mx.test_utils.assert_almost_equal(out1.asnumpy(), out2.asnumpy()) |
