219 lines
7.5 KiB
Python
219 lines
7.5 KiB
Python
import numpy as np
|
|
import pycuda.autoinit # noqa
|
|
import pycuda.driver as cuda
|
|
import tensorrt as trt
|
|
import torch
|
|
from numpy.testing import assert_array_almost_equal
|
|
|
|
INPUT_BLOB_NAME = 'input'
|
|
OUTPUT_BLOB_NAME = 'output'
|
|
|
|
|
|
def shift_mit(x, num_segments, shift_div=8):
|
|
"""Official temporal shift module.
|
|
|
|
Code Reference: https://github.com/mit-han-lab/temporal-shift-module/blob/master/ops/temporal_shift.py # noqa
|
|
Cannot convert to ONNX Model.
|
|
"""
|
|
nt, c, h, w = x.size()
|
|
n_batch = nt // num_segments
|
|
x = x.view(n_batch, num_segments, c, h, w)
|
|
|
|
fold = c // shift_div
|
|
|
|
out = torch.zeros_like(x)
|
|
out[:, :-1, :fold] = x[:, 1:, :fold] # shift left
|
|
out[:, 1:, fold:2 * fold] = x[:, :-1, fold:2 * fold] # shift right
|
|
out[:, :, 2 * fold:] = x[:, :, 2 * fold:] # not shift
|
|
|
|
return out.view(nt, c, h, w)
|
|
|
|
|
|
def shift_mmaction2(x, num_segments, shift_div=8):
|
|
"""MMAction2 temporal shift module.
|
|
|
|
Code Reference: https://github.com/open-mmlab/mmaction2/blob/master/mmaction/models/backbones/resnet_tsm.py # noqa
|
|
Could convert to ONNX Model.
|
|
"""
|
|
# [N, C, H, W]
|
|
n, c, h, w = x.size()
|
|
|
|
# [N // num_segments, num_segments, C, H*W]
|
|
# can't use 5 dimensional array on PPL2D backend for caffe
|
|
x = x.view(-1, num_segments, c, h * w)
|
|
|
|
# get shift fold
|
|
fold = c // shift_div
|
|
|
|
# split c channel into three parts:
|
|
# left_split, mid_split, right_split
|
|
left_split = x[:, :, :fold, :]
|
|
mid_split = x[:, :, fold:2 * fold, :]
|
|
right_split = x[:, :, 2 * fold:, :]
|
|
|
|
# can't use torch.zeros(*A.shape) or torch.zeros_like(A)
|
|
# because array on caffe inference must be got by computing
|
|
|
|
# shift left on num_segments channel in `left_split`
|
|
zeros = left_split - left_split
|
|
blank = zeros[:, :1, :, :]
|
|
left_split = left_split[:, 1:, :, :]
|
|
left_split = torch.cat((left_split, blank), 1)
|
|
|
|
# shift right on num_segments channel in `mid_split`
|
|
zeros = mid_split - mid_split
|
|
blank = zeros[:, :1, :, :]
|
|
mid_split = mid_split[:, :-1, :, :]
|
|
mid_split = torch.cat((blank, mid_split), 1)
|
|
|
|
# right_split: no shift
|
|
|
|
# concatenate
|
|
out = torch.cat((left_split, mid_split, right_split), 2)
|
|
|
|
# [N, C, H, W]
|
|
# restore the original dimension
|
|
return out.view(n, c, h, w)
|
|
|
|
|
|
def _tensorrt_shift_module(network,
|
|
input,
|
|
num_segments=8,
|
|
shift_div=8,
|
|
input_shape=(16, 64, 32, 32)):
|
|
"""Temporal shift module implemented by TensorRT Network Definition API."""
|
|
fold = input_shape[1] // shift_div
|
|
batch_size = input_shape[0] // num_segments
|
|
|
|
# reshape
|
|
reshape = network.add_shuffle(input)
|
|
assert reshape
|
|
reshape.reshape_dims = (batch_size, num_segments) + tuple(input_shape[-3:])
|
|
|
|
# left
|
|
left_split = network.add_slice(reshape.get_output(0),
|
|
start=(0, 1, 0, 0, 0),
|
|
shape=(batch_size, num_segments - 1, fold,
|
|
input_shape[2], input_shape[3]),
|
|
stride=(1, 1, 1, 1, 1))
|
|
assert left_split
|
|
left_split_shape = (batch_size, 1, fold, input_shape[2], input_shape[3])
|
|
left_blank = network.add_constant(shape=left_split_shape,
|
|
weights=np.zeros(left_split_shape,
|
|
np.float32))
|
|
assert left_blank
|
|
left = network.add_concatenation(
|
|
[left_split.get_output(0),
|
|
left_blank.get_output(0)])
|
|
assert left
|
|
left.axis = 1
|
|
|
|
# mid
|
|
mid_split_shape = (batch_size, 1, fold, input_shape[2], input_shape[3])
|
|
mid_blank = network.add_constant(shape=mid_split_shape,
|
|
weights=np.zeros(mid_split_shape,
|
|
np.float32))
|
|
assert mid_blank
|
|
mid_split = network.add_slice(reshape.get_output(0),
|
|
start=(0, 0, fold, 0, 0),
|
|
shape=(batch_size, num_segments - 1, fold,
|
|
input_shape[2], input_shape[3]),
|
|
stride=(1, 1, 1, 1, 1))
|
|
assert mid_split
|
|
mid = network.add_concatenation(
|
|
[mid_blank.get_output(0),
|
|
mid_split.get_output(0)])
|
|
assert mid
|
|
mid.axis = 1
|
|
|
|
# right
|
|
right = network.add_slice(reshape.get_output(0),
|
|
start=(0, 0, 2 * fold, 0, 0),
|
|
shape=(batch_size, num_segments,
|
|
input_shape[1] - 2 * fold, input_shape[2],
|
|
input_shape[3]),
|
|
stride=(1, 1, 1, 1, 1))
|
|
|
|
# concat
|
|
concat = network.add_concatenation(
|
|
[left.get_output(0),
|
|
mid.get_output(0),
|
|
right.get_output(0)])
|
|
assert concat
|
|
concat.axis = 2
|
|
|
|
# reshape
|
|
reshape2 = network.add_shuffle(concat.get_output(0))
|
|
assert reshape2
|
|
reshape2.reshape_dims = input_shape
|
|
return reshape2
|
|
|
|
|
|
def shift_tensorrt(x, num_segments, shift_div, input_shape):
|
|
"""Test TensorRT temporal shift module."""
|
|
assert isinstance(x, np.ndarray)
|
|
|
|
gLogger = trt.Logger(trt.Logger.INFO)
|
|
builder = trt.Builder(gLogger)
|
|
config = builder.create_builder_config()
|
|
|
|
# create engine
|
|
explicit_flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
|
|
network = builder.create_network(explicit_flag)
|
|
input = network.add_input(INPUT_BLOB_NAME, trt.float32, input_shape)
|
|
assert input
|
|
output = _tensorrt_shift_module(network,
|
|
input,
|
|
num_segments=num_segments,
|
|
shift_div=shift_div,
|
|
input_shape=input_shape)
|
|
assert output
|
|
|
|
# generate engine by builder/network/config
|
|
output.get_output(0).name = OUTPUT_BLOB_NAME
|
|
network.mark_output(output.get_output(0))
|
|
builder.max_batch_size = 1
|
|
builder.max_workspace_size = 1 << 20
|
|
engine = builder.build_engine(network, config)
|
|
del network
|
|
assert engine.num_bindings == 2, f'{engine.num_bindings}'
|
|
context = engine.create_execution_context()
|
|
|
|
# buffer
|
|
host_in = cuda.pagelocked_empty(trt.volume(input_shape), dtype=np.float32)
|
|
np.copyto(host_in, x.ravel())
|
|
host_out = cuda.pagelocked_empty(trt.volume(input_shape), dtype=np.float32)
|
|
devide_in = cuda.mem_alloc(host_in.nbytes)
|
|
devide_out = cuda.mem_alloc(host_out.nbytes)
|
|
bindings = [int(devide_in), int(devide_out)]
|
|
stream = cuda.Stream()
|
|
|
|
# do inference
|
|
cuda.memcpy_htod_async(devide_in, host_in, stream)
|
|
context.execute_async(bindings=bindings, stream_handle=stream.handle)
|
|
cuda.memcpy_dtoh_async(host_out, devide_out, stream)
|
|
stream.synchronize()
|
|
|
|
return np.array(host_out.reshape(*input_shape))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
INPUT_SHAPE = (16, 64, 32, 32)
|
|
assert len(INPUT_SHAPE) == 4
|
|
NUM_SEGMENTS = 8
|
|
SHIFT_DIV = 8
|
|
|
|
# inference
|
|
inputs = np.random.rand(*INPUT_SHAPE).astype(np.float32)
|
|
inputs_pytorch = torch.tensor(inputs)
|
|
with torch.no_grad():
|
|
rmit = shift_mit(inputs_pytorch, NUM_SEGMENTS, SHIFT_DIV).numpy()
|
|
rmmaction2 = shift_mmaction2(inputs_pytorch, NUM_SEGMENTS,
|
|
SHIFT_DIV).numpy()
|
|
rtensorrt = shift_tensorrt(inputs, NUM_SEGMENTS, SHIFT_DIV, INPUT_SHAPE)
|
|
|
|
# test results
|
|
assert_array_almost_equal(rmit, rtensorrt)
|
|
assert_array_almost_equal(rmmaction2, rtensorrt)
|
|
print("Tests PASSED")
|