kevinwang676 commited on
Commit
d978177
1 Parent(s): f9f60ce

Update monotonic_align/__init__.py

Browse files
Files changed (1) hide show
  1. monotonic_align/__init__.py +16 -15
monotonic_align/__init__.py CHANGED
@@ -1,19 +1,20 @@
1
- import numpy as np
2
- import torch
3
- from .monotonic_align.core import maximum_path_c
 
4
 
5
 
6
  def maximum_path(neg_cent, mask):
7
- """ Cython optimized version.
8
- neg_cent: [b, t_t, t_s]
9
- mask: [b, t_t, t_s]
10
- """
11
- device = neg_cent.device
12
- dtype = neg_cent.dtype
13
- neg_cent = neg_cent.data.cpu().numpy().astype(np.float32)
14
- path = np.zeros(neg_cent.shape, dtype=np.int32)
15
 
16
- t_t_max = mask.sum(1)[:, 0].data.cpu().numpy().astype(np.int32)
17
- t_s_max = mask.sum(2)[:, 0].data.cpu().numpy().astype(np.int32)
18
- maximum_path_c(path, neg_cent, t_t_max, t_s_max)
19
- return torch.from_numpy(path).to(device=device, dtype=dtype)
 
1
+ from numpy import zeros, int32, float32
2
+ from torch import from_numpy
3
+
4
+ from .core import maximum_path_jit
5
 
6
 
7
  def maximum_path(neg_cent, mask):
8
+ """ numba optimized version.
9
+ neg_cent: [b, t_t, t_s]
10
+ mask: [b, t_t, t_s]
11
+ """
12
+ device = neg_cent.device
13
+ dtype = neg_cent.dtype
14
+ neg_cent = neg_cent.data.cpu().numpy().astype(float32)
15
+ path = zeros(neg_cent.shape, dtype=int32)
16
 
17
+ t_t_max = mask.sum(1)[:, 0].data.cpu().numpy().astype(int32)
18
+ t_s_max = mask.sum(2)[:, 0].data.cpu().numpy().astype(int32)
19
+ maximum_path_jit(path, neg_cent, t_t_max, t_s_max)
20
+ return from_numpy(path).to(device=device, dtype=dtype)