背景
ImageNet数据集有大约1.3e6张图片,不禁令人好奇,其中的图片的分辨率的分布是怎样的呢?在机器学习的训练中,如果我们提前缩小图片的分辨率会不会大幅提升训练速度呢?
FFProbe
ffprobe是著名的多媒体工具FFmpeg的一个组件,简单地通过命令行调用即可快速查询各种多媒体文件的信息。虽然它主要的用途是查询视频,但也可以用于查询图片的信息,其中也包括了的分辨率。以下是使用的方法:
ffprobe -loglevel warning -show_entries stream=height,width -print_format json /path/to/image.jpg
它的输出大概是这样的,通过解析JSON就能获得图片的分辨率了。
{
"programs": [
],
"streams": [
{
"width": 480,
"height": 320
}
]
}
那剩下的问题就是,怎么快速调用ffprobe 1.3e6次,并统计它的结果了。
Python subprocess:20次/s
要说现在机器学习社区什么编程语言最火,那自然是Python了。这次任务也用Python来试试。最简单直接的方法,首先,用pathlib.Path.rglob()
获取一个所有图片的列表,然后逐个图片用subprocess.run
调用ffprobe,并解析返回的JSON,写到一个dict
里。这里还直接用到了之前写的用asyncio
协程来同时起多个子进程的方法。其主体部分大概是这样:
proc = await asyncio.create_subprocess_exec(
'ffprobe', '-loglevel', 'warning',
'-show_entries', 'stream=height,width',
'-print_format', 'json',
str(img),
stdout=asyncio.subprocess.PIPE, stderr=None, stdin=asyncio.subprocess.DEVNULL
)
stdout, stderr = await proc.communicate()
if proc.returncode != 0:
logger.error('ffprobe %s failed with return code %d', img, proc.returncode)
return (str(img), None, None)
probe_result = json.loads(stdout)['streams'][0]
return (str(img), probe_result['width'], probe_result['height'])
这个版本的代码没有留下来了,但最终的速度,大约只有不到20次/s,也就是说,跑完整个数据集大概要20小时,这个速度显然不能接受。但首先,我们来用cprofile看看它为啥这么慢:
ncalls tottime percall cumtime percall filename:lineno(function)
980 40.736 0.042 40.736 0.042 {built-in method _posixsubprocess.fork_exec}
4897 18.612 0.004 18.612 0.004 {built-in method posix.read}
980 0.229 0.000 59.740 0.061 .../python3.7/subprocess.py:1412(_execute_child)
绝大多数时间都花费在了内置函数fork_exec
上。究竟为啥还尚且不知。
multiprocessing + subprocess: 200次/s
加速最简单的方法,那自然是并行化了。得益于我们20核心的高性能服务器,这样的策略一般都能取得很不错的效果。最简单的方案,用multiprocessing.Process
运行多个进程,用queue.Queue
在进程间传递需要探测的图片路径和探测结果。这样的话,性能大约提升了10倍,花费了约2小时完成了这项工作。但其实我还是不太满意这样的性能。毕竟,为啥这么常用的subprocess会这么慢呢。
forkserver: 2000次/s
这事挺神奇的,但想通了的话却也在情理之中。
首先是来自一个发现,我的程序处理的图片总数量越少,处理速度就越快,并且这个效果非常明显。这挺奇怪的,处理的总数量和处理速度应该没有关系的吧,毕竟不管有多少图片,都是一张一张处理的呀。在处理图片数量少的时候,cprofile的结果大概是这样:
ncalls tottime percall cumtime percall filename:lineno(function)
20738 25.725 0.001 25.725 0.001 {built-in method posix.read}
4148 21.226 0.005 21.226 0.005 {built-in method _posixsubprocess.fork_exec}
可以看到,fork_exec的percall(每次调用)时间大概少了一个数量级。CPU利用率统计结果来看,system占用的比例也明显减少了。
在我冥思苦想了几十分钟后,终于想到了答案。这事还得从subprocess的底层实现说起。在Linux系统中,要启动一个子进程,通常需要使用fork
和execve
这两个系统调用,其中的fork
就是罪魁祸首了。fork
的作用是创建一个和当前进程一模一样的进程,虽然实际的内存页是有“写时复制”的,并不需要在fork时进行复制,但系统还要处理页表之类的数据结构,依然需要消耗一些时间。而当把整个数据集加载进内存后,整个Python进程大约会占用800M的内存,虚拟地址空间更是到了1G以上。这些占用的内存所需数据结构的维护正是导致fork
性能低的原因。而当图片总数越少时,Python进程的内存消耗就更少,自然fork得更快了。
在知道了这一点以后,解决方法就很简单了。我们使用multiprocessing.get_context('forkserver')
来获得一个context,之后用这个context来启动worker进程。这样的话,worker进程是从0开始执行的,而不是从主进程fork出来的,它的内存占用在稳定时只有18M,自然fork的速度非常快了。它的cprofile结果大概是这样的:
ncalls tottime percall cumtime percall filename:lineno(function)
416786 347.520 0.001 347.520 0.001 {built-in method posix.read}
83167 89.737 0.001 89.737 0.001 {built-in method _posixsubprocess.fork_exec}
fork_exec的速度变得非常快了。整个数据集也大概只需要12分钟就能跑完。但还能更快!
PyAV:20000次/s
但是说到底,不断起子进程本来就不是个高效的做法。进程启动,进程间通信等造成的开销很大。FFmpeg项目有个很著名的Python的绑定,叫PyAV,这个项目可以帮助你直接在进程内调用FFmpeg的功能。经过一番摸索,用PyAV来探测图片的分辨率非常的简单,甚至比启动子进程更简单:
import av
container = av.open(str(img))
s = container.streams[0]
width = s.codec_context.width
height = s.codec_context.height
直接用这个替换掉原来启动子进程的代码,效果立竿见影,又快了10倍,只需要1-2分钟就能完成。要知道,遍历一遍数据集的目录获取文件列表都需要几十秒的时间。
concurrent.futures:28092次/s
最后的一些优化了,之前为了配合我用asyncio
来启动子进程的做法,我自己用multiprocessing
实现了多进程。但这是比较底层的做法。现在既然我抛弃了子进程,那我就可以使用一些更高级的封装。通过concurrent.futures.ProcessPoolExecutor
我可以很方便地把任务分配到多个进程里。得益于更高效的实现,以及它将任务打包到更大的块中再分配给子进程的策略,性能又提高了约40%
with ProcessPoolExecutor(max_workers=WORKER) as executor:
map_iter = executor.map(probe, pending_imgs, chunksize=256)
for img_path, width, height in tqdm(map_iter, total=len(pending_imgs)):
results[img_path] = [width, height]
总结
这真可谓是个经典的例子,不要以为现在计算机的硬件已经很快了,写软件的时候就能为所欲为。通过我对软件的知识,理解和探索,我用了几个小时的时间,把程序的性能从最naïve的版本,提升了3个数量级。这能帮助以后我们能更快地完成这类任务,也让我对相关软件的认识更近一步了。
它还可以更快吗?当然可以,但是,现在这个速度对于我们来说,已经远远超过了“够用”的程度了吧。
对了,回到最初的问题,ImageNet的图像分辨率的分布究竟是怎样的呢。看看我的最终结果吧。总的来说,绝大多数图片都不是很大,提前缩小图片的分辨率应该不会对训练速度带来什么帮助。
第一次发文章,请多包涵。