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")