List转tensor(pytorch）的各种方式速度对比

List转tensor的各种方式速度对比

最近跑代码发现用GPU（还是4090!）比CPU还慢，找了各种方法都没找到问题，最后发现是在CPU和GPU之间的数据转换（numpy/list->tensor）太慢的原因！

下面是实验验证过程，不想看实验过程可以直接看结论。

1、list里面是纯数字

1.1 直接转tensor

import numpy as np
import torch
import time

l = [1 for i in range(50000000)]  # 五千万
stime = time.time()
torch.tensor(l)
etime = time.time()
print(f'直接转tensor耗时: {etime-stime}s')

直接转tensor耗时: 2.511237144470215s

1.2 先转ndarray再转tensor

l2 = [1 for i in range(50000000)]  # 五千万
stime = time.time()
l2 = np.array(l2)
etime = time.time()
print(f'先转numpy耗时: {etime-stime}s')
stime = time.time()
torch.tensor(l2)
etime = time.time()
print(f'再转tensor耗时: {etime-stime}s')

先转numpy耗时: 2.9608983993530273s
再转tensor耗时: 0.03789925575256348s

结论：List里面是纯数字时，直接转tensor快一些，但相差不大

2、list里面有numpy数组

2.1 直接转tensor

l = [np.ones(1) for i in range(50000000)]  # 五千万
stime = time.time()
torch.tensor(l)
etime = time.time()
print(f'直接转tensor耗时: {etime-stime}s')

提示：UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor.

直接转tensor耗时: 34.85922598838806s

2.2 先转ndarray再转tensor

l2 = [np.ones(1) for i in range(50000000)]  # 五千万
stime = time.time()
l2 = np.array(l2)
etime = time.time()
print(f'先转numpy耗时: {etime-stime}s')
stime = time.time()
torch.tensor(l2)
etime = time.time()
print(f'再转tensor耗时: {etime-stime}s')

先转numpy耗时: 23.368181943893433s
再转tensor耗时: 0.07274556159973145s

结论：List里面包含numpy数组时，直接转成tensor会非常慢，最好是先转成numpy之后再转tensor

3、list里面有List

3.1 直接转tensor

l = [[1] for i in range(50000000)]  # 五千万
stime = time.time()
torch.tensor(l)
etime = time.time()
print(f'直接转tensor耗时: {etime-stime}s')

直接转tensor耗时: 9.082918882369995s

3.2 先转ndarray再转tensor

l2 = [[1] for i in range(50000000)]  # 五千万
stime = time.time()
l2 = np.array(l2)
etime = time.time()
print(f'先转numpy耗时: {etime-stime}s')
stime = time.time()
torch.tensor(l2)
etime = time.time()
print(f'再转tensor耗时: {etime-stime}s')

先转numpy耗时: 37.20958995819092s
再转tensor耗时: 0.037840843200683594s

结论：List里面包含List时，就不要转numpy数组了，速度也非常慢

总结

• List里面没有numpy数组时（包括纯数字或List），直接转tensor最快
• List里面含有numpy数组时，直接转tensor非常慢，建议先转为numpy数组再转tensor
• 如果可以选择的话，List里面包含List 相比 numpy数组里面包含numpy数组，转tensor的速度更快