tensorflow源码分析之设备(device)

259 阅读1分钟

本文已参与「新人创作礼」活动,一起开启掘金创作之路。

 tensorflow支持多种设备,如GPU,CPU, TPU等待。设备其实就是计算设备。用来执行计算的。在tensorflow中,计算被封装到OP中. OP的实现Kernel实现具体计算。计算还需要数据,因此需要能把数据复制到设备中,然后进行计算。

设备属性通过proto定义了设备的属性。设备属性可以在OP中获得。从OpKernelContext中获取设备,然后可以获得设备属性。

设备上下文:DeviceContext包含这么几个主要功能,提供stream_executor, 提供CPU到非CPU(GPU)设备的Copy数据,提供本设备内部数据拷贝,提供设备内存分配Allocator

设备基类 设备依赖env,设备中有cpu线程池, 也有其他设备上的线程池,eigen tensorflow线程池,设置Allocator以分配内存给Tensor,从Proto反序化出一个tensor到设备上,维护加速设备信息:GPU/TPU

设备才是本文的重点,主要有这么几个字段,设备名,设备类型:从设备属性中获得,同步异步计算:最核心的功能,设备就是用来执行op kernel的,持有ResourceMgr,持有Env,Sync: 计算任务在设备上排队,调用此函数等执行完