转载来自:
https://blog.csdn.net/seven08290/article/details/103823364
---------------------------------------------------------------------------------------------------------
近期本人在做OCR的时候,代码本地训练loss每一次迭代都有值。
但是到服务器,每一次迭代loss都为Nan。
后来发现,服务器是有GPU 有cuda,代码要略加修改:
---- 在计算CTC loss前,将length和text加上cuda即可。 问题解决!!!如下:
preds = crnn(image)
batch_size = image.size(0)
text, length = converter.encode(label)
if str(device) == 'cpu':
# print('cpu!!!!!!!!!!!!!!!!!!!!!')
pass
else:
print('gpu!!!!!!!!!!!!!!!!!!!!!')
# 服务器 GPU训练加上cuda 否则loss一直为NaN
text = text.to(torch.long).cuda()
length = length.to(torch.long).cuda()
preds_size = Variable(torch.IntTensor([preds.size(0)] * batch_size))
cost = criterion(preds, text, preds_size, length) / batch_size