TensorFlow的图切割模块——Graph Partitioner

 

背景

[作者:

第一步——分析构建Control Flow相关信息

复制代码
 1 GraphInfo g_info;  2 if (!opts.control_flow_added) {  3   // Add the "code" for distributed execution of control flow. Code is  4   // added only for the frames that are placed on multiple devices. The  5   // new graph is an equivalent transformation of the original graph and  6   // has the property that it can be subsequently partitioned arbitrarily  7   // (down to the level of individual device) for distributed execution. 8   status = AddControlFlow(opts, g, &g_info);  9   if (!status.ok()) return status; 10 }
复制代码

第二步——构建Op的Input和Output Memory类型信息

在介绍这个过程之前,首先需要明确两种概念,他们是DeviceMemory和HostMemory。前者指的是计算设备的Memory类型,后者指的是CPU的Memory类型,它们在TensorFlow中被定义为Enum类型,代码如下所示。
复制代码
1 // MemoryType is used to describe whether input or output Tensors of 2 // an OpKernel should reside in "Host memory" (e.g., CPU memory) or 3 // "Device" Memory (CPU memory for CPU devices, GPU memory for GPU 4 // devices).5 enum MemoryType { 6   DEVICE_MEMORY = 0, 7   HOST_MEMORY = 1, 8 };
复制代码
对Op的Input和Output Memory信息进行检索并构建缓存的函数是BuildMemoryDeviceInfo,该过程构建的信息对后面真正做图切割非常重要。因为TensorFlow的Op在注册时需要不但需要指定其在各个Device上的实现版本(比如CPU版本的Op和GPU版本的Op都是分别注册到系统中的),还需要指出其Input和Output Tensor的类型以及所使用的Memory类型,即使某个Op存在GPU上的实现,它的GPU版本也有可能需要在CPU上读入数据或输出结果。例如,GPU版本的Reshape Op注册代码如下。
复制代码
 1 #define REGISTER_GPU_KERNEL(type)                               \ 2   REGISTER_KERNEL_BUILDER(Name("Reshape")                       \  3                               .Device(DEVICE_GPU)               \  4                               .HostMemory("shape")              \  5                               .TypeConstraint<type>("
                        
关键字:
50000+
5万行代码练就真实本领
17年
创办于2008年老牌培训机构
1000+
合作企业
98%
就业率

联系我们

电话咨询

0532-85025005

扫码添加微信