File size: 2,482 Bytes
b11387a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
import torch
import axengine as axe
import numpy as np
import math
def _get_source_mask(src_seq, valid_ratios) -> torch.Tensor:
"""Generate mask for source sequence.
Args:
src_seq (torch.Tensor): Image sequence. Shape :math:`(N, T, C)`.
valid_ratios (list[float]): The valid ratio of input image. For
example, if the width of the original image is w1 and the width
after padding is w2, then valid_ratio = w1/w2. Source mask is
used to cover the area of the padding region.
Returns:
Tensor or None: Source mask. Shape :math:`(N, T)`. The region of
padding area are False, and the rest are True.
"""
N, T, _ = src_seq.size()
mask = None
if len(valid_ratios) > 0:
mask = src_seq.new_zeros((N, T), device=src_seq.device)
for i, valid_ratio in enumerate(valid_ratios):
valid_width = min(T, math.ceil(T * valid_ratio))
mask[i, :valid_width] = 1
return mask
onnx_bb_encoder = axe.InferenceSession("backbone_encoder.axmodel")
onnx_decoder = axe.InferenceSession("decoder.axmodel")
input_image = torch.tensor(np.load('input_tensor/input_image.npy'))
out_enc = onnx_bb_encoder.run(["output"], {"input": np.array(input_image.cpu())})[0]
out_enc = torch.tensor(out_enc)
data_samples = None
N = out_enc.size(0)
init_target_seq = torch.tensor(np.load('input_tensor/init_target_seq.npy')).to(torch.int32)
outputs = []
max_seq_len = 25
for step in range(0, max_seq_len):
valid_ratios = [1.0 for _ in range(out_enc.size(0))]
if data_samples is not None:
valid_ratios = []
for data_sample in data_samples:
valid_ratios.append(data_sample.get('valid_ratio'))
src_mask = _get_source_mask(out_enc, valid_ratios)
# step_result = model_decoder(init_target_seq,out_enc,src_mask,step)
step_result = onnx_decoder.run(["output"],{'init_target_seq':np.array(init_target_seq),
'out_enc':np.array(out_enc),
'src_mask':np.array(src_mask),
'step':np.array([step]).astype(np.int32)})[0][0]
step_result = torch.tensor(step_result)
outputs.append(step_result)
_, step_max_index = torch.max(step_result, dim=-1)
init_target_seq[:, step + 1] = step_max_index
outputs = torch.stack(outputs, dim=1)
np.save('output_tensor/outputs.npy',outputs)
|