| | |
|---|
| | |
|
|---|
| | | return mask.type(dtype).to(device) if device is not None else mask.type(dtype)
|
|---|
| | |
|
|---|
| | | def sequence_mask_scripts(lengths, maxlen:int):
|
|---|
| | | row_vector = torch.arange(0, maxlen, 1).type(lengths.dtype).to(lengths.device)
|
|---|
| | | matrix = torch.unsqueeze(lengths, dim=-1)
|
|---|
| | | mask = row_vector < matrix
|
|---|
| | | return mask.type(torch.float32).to(lengths.device)
|
|---|
| | |
|
|---|
| | | class CifPredictorV2(nn.Module):
|
|---|
| | | def __init__(self, model):
|
|---|
| | |
|---|
| | |
|
|---|
| | | return hidden, alphas, token_num_floor
|
|---|
| | |
|
|---|
| | |
|
|---|
| | | # @torch.jit.script
|
|---|
| | | # def cif(hidden, alphas, threshold: float):
|
|---|
| | | # batch_size, len_time, hidden_size = hidden.size()
|
|---|
| | | # threshold = torch.tensor([threshold], dtype=alphas.dtype).to(alphas.device)
|
|---|
| | | #
|
|---|
| | | # # loop varss
|
|---|
| | | # integrate = torch.zeros([batch_size], device=hidden.device)
|
|---|
| | | # frame = torch.zeros([batch_size, hidden_size], device=hidden.device)
|
|---|
| | | # # intermediate vars along time
|
|---|
| | | # list_fires = []
|
|---|
| | | # list_frames = []
|
|---|
| | | #
|
|---|
| | | # for t in range(len_time):
|
|---|
| | | # alpha = alphas[:, t]
|
|---|
| | | # distribution_completion = torch.ones([batch_size], device=hidden.device) - integrate
|
|---|
| | | #
|
|---|
| | | # integrate += alpha
|
|---|
| | | # list_fires.append(integrate)
|
|---|
| | | #
|
|---|
| | | # fire_place = integrate >= threshold
|
|---|
| | | # integrate = torch.where(fire_place,
|
|---|
| | | # integrate - torch.ones([batch_size], device=hidden.device),
|
|---|
| | | # integrate)
|
|---|
| | | # cur = torch.where(fire_place,
|
|---|
| | | # distribution_completion,
|
|---|
| | | # alpha)
|
|---|
| | | # remainds = alpha - cur
|
|---|
| | | #
|
|---|
| | | # frame += cur[:, None] * hidden[:, t, :]
|
|---|
| | | # list_frames.append(frame)
|
|---|
| | | # frame = torch.where(fire_place[:, None].repeat(1, hidden_size),
|
|---|
| | | # remainds[:, None] * hidden[:, t, :],
|
|---|
| | | # frame)
|
|---|
| | | #
|
|---|
| | | # fires = torch.stack(list_fires, 1)
|
|---|
| | | # frames = torch.stack(list_frames, 1)
|
|---|
| | | # list_ls = []
|
|---|
| | | # len_labels = torch.floor(alphas.sum(-1)).int()
|
|---|
| | | # max_label_len = len_labels.max()
|
|---|
| | | # for b in range(batch_size):
|
|---|
| | | # fire = fires[b, :]
|
|---|
| | | # l = torch.index_select(frames[b, :, :], 0, torch.nonzero(fire >= threshold).squeeze())
|
|---|
| | | # pad_l = torch.zeros([int(max_label_len - l.size(0)), int(hidden_size)], device=hidden.device)
|
|---|
| | | # list_ls.append(torch.cat([l, pad_l], 0))
|
|---|
| | | # return torch.stack(list_ls, 0), fires
|
|---|
| | |
|
|---|
| | |
|
|---|
| | | @torch.jit.script
|
|---|
| | | def cif(hidden, alphas, threshold: float):
|
|---|
| | | batch_size, len_time, hidden_size = hidden.size()
|
|---|
| | | threshold = torch.tensor([threshold], dtype=alphas.dtype).to(alphas.device)
|
|---|
| | |
|
|---|
| | | # loop varss
|
|---|
| | | integrate = torch.zeros([batch_size], device=hidden.device)
|
|---|
| | | frame = torch.zeros([batch_size, hidden_size], device=hidden.device)
|
|---|
| | | integrate = torch.zeros([batch_size], dtype=alphas.dtype, device=hidden.device)
|
|---|
| | | frame = torch.zeros([batch_size, hidden_size], dtype=hidden.dtype, device=hidden.device)
|
|---|
| | | # intermediate vars along time
|
|---|
| | | list_fires = []
|
|---|
| | | list_frames = []
|
|---|
| | |
|
|---|
| | | for t in range(len_time):
|
|---|
| | | alpha = alphas[:, t]
|
|---|
| | | distribution_completion = torch.ones([batch_size], device=hidden.device) - integrate
|
|---|
| | | distribution_completion = torch.ones([batch_size], dtype=alphas.dtype, device=hidden.device) - integrate
|
|---|
| | |
|
|---|
| | | integrate += alpha
|
|---|
| | | list_fires.append(integrate)
|
|---|
| | |
|
|---|
| | | fire_place = integrate >= threshold
|
|---|
| | | integrate = torch.where(fire_place,
|
|---|
| | | integrate - torch.ones([batch_size], device=hidden.device),
|
|---|
| | | integrate - torch.ones([batch_size], dtype=alphas.dtype, device=hidden.device),
|
|---|
| | | integrate)
|
|---|
| | | cur = torch.where(fire_place,
|
|---|
| | | distribution_completion,
|
|---|
| | |
|---|
| | |
|
|---|
| | | fires = torch.stack(list_fires, 1)
|
|---|
| | | frames = torch.stack(list_frames, 1)
|
|---|
| | | list_ls = []
|
|---|
| | | len_labels = torch.round(alphas.sum(-1)).int()
|
|---|
| | | max_label_len = len_labels.max()
|
|---|
| | |
|
|---|
| | | fire_idxs = fires >= threshold
|
|---|
| | | frame_fires = torch.zeros_like(hidden)
|
|---|
| | | max_label_len = frames[0, fire_idxs[0]].size(0)
|
|---|
| | | for b in range(batch_size):
|
|---|
| | | fire = fires[b, :]
|
|---|
| | | l = torch.index_select(frames[b, :, :], 0, torch.nonzero(fire >= threshold).squeeze())
|
|---|
| | | pad_l = torch.zeros([int(max_label_len - l.size(0)), int(hidden_size)], device=hidden.device)
|
|---|
| | | list_ls.append(torch.cat([l, pad_l], 0))
|
|---|
| | | return torch.stack(list_ls, 0), fires
|
|---|
| | |
|
|---|
| | |
|
|---|
| | | def CifPredictorV2_test():
|
|---|
| | | x = torch.rand([2, 21, 2])
|
|---|
| | | x_len = torch.IntTensor([6, 21])
|
|---|
| | | frame_fire = frames[b, fire_idxs[b]]
|
|---|
| | | frame_len = frame_fire.size(0)
|
|---|
| | | frame_fires[b, :frame_len, :] = frame_fire
|
|---|
| | |
|
|---|
| | | mask = sequence_mask(x_len, maxlen=x.size(1), dtype=x.dtype)
|
|---|
| | | x = x * mask[:, :, None]
|
|---|
| | | |
|---|
| | | predictor_scripts = torch.jit.script(CifPredictorV2(2, 1, 1))
|
|---|
| | | # cif_output, cif_length, alphas, cif_peak = predictor_scripts(x, mask=mask[:, None, :])
|
|---|
| | | predictor_scripts.save('test.pt')
|
|---|
| | | loaded = torch.jit.load('test.pt')
|
|---|
| | | cif_output, cif_length, alphas, cif_peak = loaded(x, mask=mask[:, None, :])
|
|---|
| | | # print(cif_output)
|
|---|
| | | print(predictor_scripts.code)
|
|---|
| | | # predictor = CifPredictorV2(2, 1, 1)
|
|---|
| | | # cif_output, cif_length, alphas, cif_peak = predictor(x, mask=mask[:, None, :])
|
|---|
| | | print(cif_output)
|
|---|
| | |
|
|---|
| | |
|
|---|
| | | def CifPredictorV2_export_test():
|
|---|
| | | x = torch.rand([2, 21, 2])
|
|---|
| | | x_len = torch.IntTensor([6, 21])
|
|---|
| | | |
|---|
| | | mask = sequence_mask(x_len, maxlen=x.size(1), dtype=x.dtype)
|
|---|
| | | x = x * mask[:, :, None]
|
|---|
| | | |
|---|
| | | # predictor_scripts = torch.jit.script(CifPredictorV2(2, 1, 1))
|
|---|
| | | # cif_output, cif_length, alphas, cif_peak = predictor_scripts(x, mask=mask[:, None, :])
|
|---|
| | | predictor = CifPredictorV2(2, 1, 1)
|
|---|
| | | predictor_trace = torch.jit.trace(predictor, (x, mask[:, None, :]))
|
|---|
| | | predictor_trace.save('test_trace.pt')
|
|---|
| | | loaded = torch.jit.load('test_trace.pt')
|
|---|
| | | |
|---|
| | | x = torch.rand([3, 30, 2])
|
|---|
| | | x_len = torch.IntTensor([6, 20, 30])
|
|---|
| | | mask = sequence_mask(x_len, maxlen=x.size(1), dtype=x.dtype)
|
|---|
| | | x = x * mask[:, :, None]
|
|---|
| | | cif_output, cif_length, alphas, cif_peak = loaded(x, mask=mask[:, None, :])
|
|---|
| | | print(cif_output)
|
|---|
| | | # print(predictor_trace.code)
|
|---|
| | | # predictor = CifPredictorV2(2, 1, 1)
|
|---|
| | | # cif_output, cif_length, alphas, cif_peak = predictor(x, mask=mask[:, None, :])
|
|---|
| | | # print(cif_output)
|
|---|
| | |
|
|---|
| | |
|
|---|
| | | if __name__ == '__main__':
|
|---|
| | | # CifPredictorV2_test()
|
|---|
| | | CifPredictorV2_export_test() |
|---|
| | | if frame_len >= max_label_len:
|
|---|
| | | max_label_len = frame_len
|
|---|
| | | frame_fires = frame_fires[:, :max_label_len, :]
|
|---|
| | | return frame_fires, fires
|
|---|
