Open lcylcy opened 2 years ago
import torch import numpy as np from warprnnt_pytorch import RNNTLoss acts = np.random.rand(2,2,3,5) labels = np.array([[1, 2],[2,2]]) act_length = np.array([2,2]) label_length = np.array([2,2]) blank = 1 reduction = 'mean' print('----------------------------------------') rnnt_loss = RNNTLoss(blank,reduction) acts_t = torch.tensor(acts,dtype=torch.float) acts_t.requires_grad = True labels_t = torch.tensor(labels,dtype=torch.int) act_length_t = torch.tensor(act_length,dtype=torch.int) label_length_t = torch.tensor(label_length,dtype=torch.int) loss = rnnt_loss(acts_t,labels_t,act_length_t,label_length_t) loss.backward() print(loss) print(acts_t.grad) print('-----------------------------------------') rnnt_loss = RNNTLoss(blank,reduction).cuda() acts_t = torch.tensor(acts,dtype=torch.float).cuda() acts_t.requires_grad = True labels_t = torch.tensor(labels,dtype=torch.int).cuda() act_length_t = torch.tensor(act_length,dtype=torch.int).cuda() label_length_t = torch.tensor(label_length,dtype=torch.int).cuda() loss = rnnt_loss(acts_t,labels_t,act_length_t,label_length_t) loss.backward() print(loss) print(acts_t.grad) print('-----------------------------------------')
---------------------------------------- tensor([4.6379], grad_fn=<_RNNTBackward>) tensor([[[[ 0.0550, -0.2406, 0.0716, 0.0589, 0.0550], [ 0.0492, -0.1081, -0.0603, 0.0419, 0.0774], [ 0.0290, -0.1254, 0.0351, 0.0291, 0.0322]], [[ 0.0346, -0.1494, 0.0226, 0.0539, 0.0383], [ 0.0726, 0.0889, -0.2694, 0.0452, 0.0627], [ 0.0730, -0.4038, 0.1470, 0.0579, 0.1259]]], [[[ 0.0749, -0.0353, -0.2435, 0.1090, 0.0948], [ 0.0529, -0.1224, -0.0426, 0.0533, 0.0587], [ 0.0290, -0.1383, 0.0383, 0.0461, 0.0249]], [[ 0.0163, 0.0226, -0.0751, 0.0122, 0.0240], [ 0.0655, 0.0462, -0.2269, 0.0494, 0.0658], [ 0.0755, -0.3531, 0.1155, 0.0999, 0.0623]]]]) ----------------------------------------- tensor([4.6379], device='cuda:0', grad_fn=<_RNNTBackward>) tensor([[[[ 0.0885, -0.3869, 0.1152, 0.0947, 0.0885], [ 0.0492, -0.1081, -0.0603, 0.0419, 0.0774], [ 0.0290, -0.1254, 0.0351, 0.0291, 0.0322]], [[ 0.0346, -0.1494, 0.0226, 0.0539, 0.0383], [ 0.0726, 0.0889, -0.2694, 0.0452, 0.0627], [ 0.0730, -0.4038, 0.1470, 0.0579, 0.1259]]], [[[ 0.0749, -0.0353, -0.2435, 0.1090, 0.0948], [ 0.0529, -0.1224, -0.0426, 0.0533, 0.0587], [ 0.0290, -0.1383, 0.0383, 0.0461, 0.0249]], [[ 0.0163, 0.0226, -0.0751, 0.0122, 0.0240], [ 0.0655, 0.0462, -0.2269, 0.0494, 0.0658], [ 0.0755, -0.3531, 0.1155, 0.0999, 0.0623]]]], device='cuda:0') -----------------------------------------
Please have a try with https://github.com/csukuangfj/optimized_transducer
It has a consistent gradient for CPU and CUDA.