Unverified Commit a377f0bd authored by functionxu123's avatar functionxu123 Committed by GitHub
Browse files

[Misc]: optimize eager mode host time (#4196)


Co-authored-by: default avatarxuhao <xuhao@cambricon.com>
parent e9d3aa04
...@@ -17,6 +17,7 @@ from typing import (Any, AsyncIterator, Awaitable, Callable, Dict, Generic, ...@@ -17,6 +17,7 @@ from typing import (Any, AsyncIterator, Awaitable, Callable, Dict, Generic,
Hashable, List, Optional, OrderedDict, Tuple, TypeVar, Hashable, List, Optional, OrderedDict, Tuple, TypeVar,
Union) Union)
import numpy as np
import psutil import psutil
import torch import torch
...@@ -501,11 +502,6 @@ def str_to_int_tuple(s: str) -> Tuple[int, ...]: ...@@ -501,11 +502,6 @@ def str_to_int_tuple(s: str) -> Tuple[int, ...]:
f"(e.g., 1, 2, 3). Given input: {s}") from e f"(e.g., 1, 2, 3). Given input: {s}") from e
def pad_to_max_length(x: List[int], max_len: int, pad: int) -> List[int]:
assert len(x) <= max_len
return x + [pad] * (max_len - len(x))
def make_tensor_with_pad( def make_tensor_with_pad(
x: List[List[int]], x: List[List[int]],
max_len: int, max_len: int,
...@@ -518,7 +514,10 @@ def make_tensor_with_pad( ...@@ -518,7 +514,10 @@ def make_tensor_with_pad(
The padding is applied to the end of each inner list until it reaches The padding is applied to the end of each inner list until it reaches
`max_len`. `max_len`.
""" """
padded_x = [pad_to_max_length(x_i, max_len, pad) for x_i in x] padded_x = np.zeros([len(x), max_len], dtype=np.int32) + pad
for ind, blocktb in enumerate(x):
assert len(blocktb) <= max_len
padded_x[ind, :len(blocktb)] = blocktb
return torch.tensor(padded_x, dtype=dtype, device=device) return torch.tensor(padded_x, dtype=dtype, device=device)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment