1, 数据要设置在gpu上存储,如:
a = nd.array([1,2,3], ctx=mx.gpu())
2,不只是第一次权重初始化的时候需要设置ctx,在后来每次调用intialize()函数的时候也要设置ctx,具体设置如下:
net.collect_params().initialize(force_reinit=True, ctx=mx.gpu())
3,iterator里出来的数据也要设置ctx。如:
for data, label in data_iter_train:
data = data.as_in_context(ctx)
label = label.as_in_context(ctx)
最开始的时候,我出现报错:
RuntimeError: Parameter sequential0_dense0_weight was not initialized on context cpu(0). It was only initialized on [gpu(0)]
然后按照上面那3点设置好之后,就不报错了。