Sequential 类的设备迁移

最新推荐文章于 2024-03-06 21:51:04 发布

weixin_30820077

最新推荐文章于 2024-03-06 21:51:04 发布

阅读量113

点赞数

文章标签：人工智能

原文链接：http://www.cnblogs.com/chenyliang/p/9493448.html

版权

之前因为RNN模块没有export 方法，直接用了 cpickle 强行保存。现在要载入保存的数据，用于inference。需要解决训练时的context和载入时 device不一致的问题。
找了下，发现ParameterDict里面有个 reset_ctx可以用:

import mxnet as mx                                                                                                                                           
import numpy as np
nn = mx.gluon.nn
net = nn.Sequential()
net.add(\
    ¦   nn.Dense(10))
ctx = mx.cpu()
_x = np.random.randint(0,256,(5,199))
x = mx.nd.array(_x)
net.initialize()


y= net(x)
print y
print('cross to gpu device...')
ctx = mx.gpu()
x = x.as_in_context( ctx )
try:
    y = net(x)
except:
    print 'forward failed, try reset_ctx for ParameterDict...'
    net.collect_params().reset_ctx( ctx )
y= net(x)
print y
print 'test ok'

转载于:https://www.cnblogs.com/chenyliang/p/9493448.html