Gluon-ts: DeepState inference on GPU

[mharvan]

Description

A DeepState model trained on GPU fails when trying to do inference because some data is on GPU and some data is on CPU. Similar training and inference code works without issues for other models (DeepAREstimator, SimpleFeedForwardEstimator,...)

When I train DeepStateEstimator on CPU only then inference works without issues.

To Reproduce

Modify examples/benchmark_m4.py:

diff --git a/examples/benchmark_m4.py b/examples/benchmark_m4.py
index 2a37b0f..b289147 100644
--- a/examples/benchmark_m4.py
+++ b/examples/benchmark_m4.py
@@ -24,46 +24,24 @@ from gluonts.distribution.piecewise_linear import PiecewiseLinearOutput
 from gluonts.evaluation import Evaluator
 from gluonts.evaluation.backtest import make_evaluation_predictions
 from gluonts.model.deepar import DeepAREstimator
+from gluonts.model.deepstate import DeepStateEstimator
 from gluonts.model.seq2seq import MQCNNEstimator
 from gluonts.model.simple_feedforward import SimpleFeedForwardEstimator
 from gluonts.trainer import Trainer

 datasets = [
     "m4_hourly",
-    "m4_daily",
-    "m4_weekly",
-    "m4_monthly",
-    "m4_quarterly",
-    "m4_yearly",
 ]

-epochs = 100
+epochs = 2
 num_batches_per_epoch = 50

 estimators = [
     partial(
-        SimpleFeedForwardEstimator,
+        DeepStateEstimator,
         trainer=Trainer(
-            epochs=epochs, num_batches_per_epoch=num_batches_per_epoch
-        ),
-    ),
-    partial(
-        DeepAREstimator,
-        trainer=Trainer(
-            epochs=epochs, num_batches_per_epoch=num_batches_per_epoch
-        ),
-    ),
-    partial(
-        DeepAREstimator,
-        distr_output=PiecewiseLinearOutput(8),
-        trainer=Trainer(
-            epochs=epochs, num_batches_per_epoch=num_batches_per_epoch
-        ),
-    ),
-    partial(
-        MQCNNEstimator,
-        trainer=Trainer(
-            epochs=epochs, num_batches_per_epoch=num_batches_per_epoch
+            epochs=epochs, num_batches_per_epoch=num_batches_per_epoch,
+            ctx='gpu'
         ),
     ),
 ]

Error Message

$ python examples/benchmark_m4.py
[...]
INFO:root:Start model training
INFO:root:Number of parameters in DeepStateTrainingNetwork: 13527
INFO:root:Epoch[0] Learning rate is 0.001
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [02:31<00:00, 3.02s/it, avg_epoch_loss=0.167]
INFO:root:Epoch[0] Elapsed time 151.200 seconds
INFO:root:Epoch[0] Evaluation metric 'epoch_loss'=0.166713
INFO:root:Epoch[1] Learning rate is 0.001
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [02:08<00:00, 2.56s/it, avg_epoch_loss=0.11]
INFO:root:Epoch[1] Elapsed time 128.203 seconds
INFO:root:Epoch[1] Evaluation metric 'epoch_loss'=0.109923
INFO:root:Loading parameters from best epoch (1)
INFO:root:Final loss: 0.10992345235077665 (occurred at epoch 1)
INFO:root:End model training
Running evaluation: 0%| | 0/414 [00:02 [15:06:21] src/imperative/./imperative_utils.h:70: Check failed: inputs[i]->ctx().dev_mask() == ctx.dev_mask() (2 vs. 1) Operator broadcast_mul require all inputs live on the same context. But the first argument is on cpu(0) while the 2-th argument is on gpu(0)

Stack trace returned 10 entries:
[bt] (0) /home/matus/anaconda3/envs/gluonts-src/lib/python3.7/site-packages/mxnet/libmxnet.so(+0x3f6cba) [0x7f797520bcba]
[bt] (1) /home/matus/anaconda3/envs/gluonts-src/lib/python3.7/site-packages/mxnet/libmxnet.so(+0x3f72e1) [0x7f797520c2e1]
[bt] (2) /home/matus/anaconda3/envs/gluonts-src/lib/python3.7/site-packages/mxnet/libmxnet.so(mxnet::imperative::GetContext(nnvm::NodeAttrs const&, std::vector > const&, std::vector > const&, mxnet::Context const&)+0x4ed) [0x7f7977c626dd]
[bt] (3) /home/matus/anaconda3/envs/gluonts-src/lib/python3.7/site-packages/mxnet/libmxnet.so(mxnet::Imperative::Invoke(mxnet::Context const&, nnvm::NodeAttrs const&, std::vector > const&, std::vector > const&)+0x2ee) [0x7f7977c69afe]
[bt] (4) /home/matus/anaconda3/envs/gluonts-src/lib/python3.7/site-packages/mxnet/libmxnet.so(+0x2d57e39) [0x7f7977b6ce39]
[bt] (5) /home/matus/anaconda3/envs/gluonts-src/lib/python3.7/site-packages/mxnet/libmxnet.so(MXImperativeInvokeEx+0x6f) [0x7f7977b6d42f]
[bt] (6) /home/matus/anaconda3/envs/gluonts-src/lib/python3.7/lib-dynload/../../libffi.so.6(ffi_call_unix64+0x4c) [0x7f79b031eec0]
[bt] (7) /home/matus/anaconda3/envs/gluonts-src/lib/python3.7/lib-dynload/../../libffi.so.6(ffi_call+0x22d) [0x7f79b031e87d]
[bt] (8) /home/matus/anaconda3/envs/gluonts-src/lib/python3.7/lib-dynload/_ctypes.cpython-37m-x86_64-linux-gnu.so(_ctypes_callproc+0x2ce) [0x7f79ee368ede]
[bt] (9) /home/matus/anaconda3/envs/gluonts-src/lib/python3.7/lib-dynload/_ctypes.cpython-37m-x86_64-linux-gnu.so(+0x12914) [0x7f79ee369914]

Environment

Ubuntu 16.04
conda 4.7.10
Python 3.7.4

$ pip list | egrep 'gluon|mxnet'
gluonts 0.3.4.dev22+g515c45c /home/matus/src/gluon-ts/src
mxnet-cu101mkl 1.4.1

$ git log | head -n 1
commit 29ab758c9ddebdbba39a12597a1d5fcaaf906a97

[lotsoffood]

I tried installing the jitter_gpu branch to fix the error, but unfortunately get the same error. Is there a way to train on the GPU (setting ctx to "GPU") but predict using the CPU?

[dcmaddix]

Yes we are still working on fixing the issue on the jitter_gpu branch. So far we have added some tests for the gpu, which are currently passing. We will expand our gpu coverage with your case. We have found numerical issues with mxnet on the gpu where the cholesky factor does not throw an error but has NaNs in its entries where there is no issue on the CPU. You could also try using higher precision np.float64.
Yes you can perform training on GPU by passing ctx=mx.Context(“gpu”) into the ‘Trainer’. Then transfer the data to the cpu and use ctx=mx.Context(“cpu”) for prediction. This is what I have done in the past and would recommend for now until we fix the issue. Thank you!

[lotsoffood]

@dcmaddix Thank you! I appreciate all the work you are all doing on this. I'll use the context setting method for now.

[lotsoffood]

I get the following error when running predictor.ctx = mx.Context("cpu") (with predictor being a DeepStateEstimator model trained on a GPU):

RuntimeError: Parameter 'deepstatepredictionnetwork1_featureembedder0_cat_0_embedding_weight' was not initialized on context cpu(0). It was only initialized on [gpu(0)].

[lotsoffood]

Same error when serializing and then deserializing with CPU context:

from path import Path
from gluonts.model.predictor import Predictor
predictor.serialize(Path('temp_model1'))
with mx.Context(mx.cpu(0)):
    predictor = Predictor.deserialize(Path('temp_model1'))

[JesseLi666]

I think this may help:
deepssm_predictor = Predictor.deserialize(Path('temp_model1'), ctx=mx.cpu())

[andmib]

I'm unable to get the Deep State estimator to predict using the GPU, even with the latest master. Does anybody have any examples that they've used that worked?

[mbrys0n]

@JesseLi666's suggestion worked for me. Try serialising a trained predictor using a Trainer with GPU context and deserializing it with a cpu context:

trainer = Trainer(ctx=mx.context.gpu(),...)
estimator = deepstate.DeepStateEstimator(...)
predictor.serialize(Path(filename))

predictor = Predictor.deserialize(Path(filename),ctx=mx.cpu(0)) # fix for deepstate

[rmitsch]

@JesseLi666's suggestion worked for me. Try serialising a trained predictor using a Trainer with GPU context and deserializing it with a cpu context:

`trainer = Trainer(ctx=mx.context.gpu(),...)
estimator = deepstate.DeepStateEstimator(...)
predictor.serialize(Path(filename))

predictor = Predictor.deserialize(Path(filename),ctx=mx.cpu(0)) # fix for deepstate
`

That's predicting without the GPU though, no? I'm interested in a solution as well. @JesseLi666's approach is a useful workaround, but I'd be thankful for a fix.