以下的讨论是基于:
MXNet版本: mxnetcu101=1.5.0
操作系统: Ubuntu 18.04
继承mxnet.gluon.loss.Loss
自定义损失函数时,涉及到一个条件运算,在多GPU上无法正常计算,不报错,一直卡在cudnn auto tune 的阶段。
def hybrid_forward(self, F, pred, label):
loss=xxxx # loss 为一个 (num_sample,) 的张量
if loss[self.idx] > self.threshold:
operation 1
else:
operation 2
self.idx
和 self.threshold
均为标量,构造损失函数的时候指定。
上述代码在CPU和单张GPU上均能正常运行,当使用多GPU并行训练的时候,就卡在语句if loss[self.idx] > self.threshold:
。不使用条件运算,多GPU时仅运行operation 1
或者operation 2
均可。
尝试过使用如下方式,仍不能正常运行,请问是什么问题呢?
def __ini__(self):
self.threshold = nd.array([threshold,], dtype='float32')
def hybrid_forward(self, F, pred, label):
loss=xxxx # loss 为一个 (num_sample,) 的张量
if F.greater(loss[self.idx], self.threshold.as_in_context(loss.context)):
operation 1
else:
operation 2