Sequential 类的设备迁移

最新推荐文章于 2025-06-23 16:34:34 发布

weixin_30820077

最新推荐文章于 2025-06-23 16:34:34 发布

阅读量134

点赞数

CC 4.0 BY-SA版权

文章标签：人工智能

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

本文介绍了一种在不同设备上迁移RNN模型参数的方法。通过使用mxnet的ParameterDict组件中的reset_ctx方法解决了训练时的context与载入时device不一致的问题。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

之前因为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