@dambo @rwightman I performed some other timing tests to try to understand if bottleneck could be in CPU to GPU transfer:

Move 6x 3-element np.arrays from CPU to GPU
Move a single 6x3 np.array from CPU to GPU
Create  6x 3-element torch.tensor directy on GPU whose elements are np.float variables
Create 6x 3-element torch.tensor directly on GPU whose elements are constant floats (hard coded)

I did these 4 tests using np.float16, np.float32 and np.float64, as well as on my desktop and a Jetson Nano. The findings are quiet interesting and counter-intuitive:

Faster operation is of a single 6x3 np.array from CPU to GPU, both on desktop and Jetson Nano. It is  around 4x faster than the rest of operations and even faster than creating 6x 3-element torch.tensor directly on GPU.
Using np.float64 is slightly faster than both np.float32 and np.float16
Transfer are 10x faster on the desktop than on Jetson Nano
Effectively, tranfering 6x 3-element np.arrays has transfer times as high as 5.5ms - 6.5ms, which is already too much for calculations at 100Hz in which the limit is 10ms.

It seems that one approach to improve performance is by transferring all data in a single np.array rather than small pieces on several individual np.arrays.
However, does this makes sense at all? Is it really possible that transferring a 6x3 np.array is faster than creating 6 torch.tensors directly on GPU? Sounds to good to be true. I understand that batch processing is one of the keypoints on using a GPU… but “batch transferring”… not really sure.
Any further comments are appreciate.
Below are the graphs:

jetsonfloat16.png1855×1022 32.2 KB
 
jetsonfloat64.png1855×1022 30.5 KB
 
jetsonfloat32.png1855×1022 33 KB
 
desktopfloat64.png1855×1022 39.3 KB
 
desktopfloat16.png1855×1022 37.7 KB
 
desktopfloat32.png1855×1022 37.8 KB