Undo hybridize to pickle model

Gluon
1

Is there a way to undo hybridization of a network in order to pickle it?

The use case is as follows:
I put together a custom graph in Dask to perform a full machine learning pipeline, e.g. data loading, cleaning, splitting, whatever. at the end of this pipeline there is the actual ML model to make predictions from the processed data. The dask pipeline can be pickled if i use sklearn or xgboost models. but of course fails when using a hybridized mxnet block.

thanks
andre

2 
net.hybridize(active=False)

should do the trick. Tip, for greater performance, when hybridizing, try also

net.hybridize(static_alloc=True,static_shape=True)

a bit more memory footprint, but important performance gain (x3-x4 with additional options).

Cheers

3

Your suggestion somehow does not solve my problem, maybe I was not clear in the description of my problem. Here minimal working example.

import pickle

import mxnet as mx
from mxnet import gluon
from mxnet.gluon import nn
from mxnet.gluon.nn import Conv2D, MaxPool2D


class SomeNet(nn.HybridBlock):
    def __init__(self):
        super(SomeNet, self).__init__()
        self.cv = Conv2D(channels=16, kernel_size=3, strides=1, padding=1)
        self.mp = MaxPool2D(pool_size=(2,2), strides=2)

    def hybrid_forward(self, F, x):
        out1 = self.cv(x)
        out2 = self.mp(out1)
        return out2


mynet = SomeNet()
mynet.collect_params().initialize(mx.init.Xavier())

# only works if next line is commented out
# mynet.hybridize(active=False)


with open('mymodel.pkl', 'wb') as picklefile:
    pickle.dump(mynet, picklefile)
4

Hi, I think the issue is that you need to do a first pass in the network, so it is initialized and then pickle/unpickle it. The following works for me:

import pickle

import mxnet as mx
from mxnet import gluon, nd
from mxnet.gluon import nn
from mxnet.gluon.nn import Conv2D, MaxPool2D

class SomeNet(nn.HybridBlock):
    def __init__(self):
        super(SomeNet, self).__init__()
        self.cv = Conv2D(channels=16, kernel_size=3, strides=1, padding=1)
        self.mp = MaxPool2D(pool_size=(2,2), strides=2)

    def hybrid_forward(self, F, x):
        out1 = self.cv(x)
        out2 = self.mp(out1)
        return out2

mynet = SomeNet()
mynet.initialize(mx.init.Xavier())
mynet.hybridize()

# This step is cruicial, to initialize the weights of the network (which are randomly initialized)
xx = nd.random.uniform(shape = [16,5,64,64])
out = mynet(xx)

mynet.hybridize(active=False)
with open('mymodel.pkl', 'wb') as picklefile:
    pickle.dump(mynet, picklefile)
with open('mymodel.pkl', 'rb') as picklefile:
    mynet2 = pickle.load(picklefile)

out2 = mynet2(xx)
out = mynet(xx)

out2-out
#prints 
[[[[0. 0. 0. ... 0. 0. 0.]
   [0. 0. 0. ... 0. 0. 0.]
   [0. 0. 0. ... 0. 0. 0.]
   ...
   [0. 0. 0. ... 0. 0. 0.]
   [0. 0. 0. ... 0. 0. 0.]
   [0. 0. 0. ... 0. 0. 0.]]

  [[0. 0. 0. ... 0. 0. 0.]
   [0. 0. 0. ... 0. 0. 0.]
   [0. 0. 0. ... 0. 0. 0.]
   ...
   [0. 0. 0. ... 0. 0. 0.]
   [0. 0. 0. ... 0. 0. 0.]
   [0. 0. 0. ... 0. 0. 0.]]

  [[0. 0. 0. ... 0. 0. 0.]
   [0. 0. 0. ... 0. 0. 0.]
   [0. 0. 0. ... 0. 0. 0.]
   ...
   [0. 0. 0. ... 0. 0. 0.]
   [0. 0. 0. ... 0. 0. 0.]
   [0. 0. 0. ... 0. 0. 0.]]

  ...
# more zeros

edit: I think the issue is the first forward pass due to deferred initialization, am not sure.

1 Like