tensorflow resource notes

在单机版中,创建session的时候会创建devices,创建device的时候会为每个device创建一个ResourceMgr。每个ResourceMgr可以有多个container。所以,在单机版中,session被销毁后它所对应的ResourceMgr里的那些resource也会被销毁。但是分布式版中就不是如此了。

ResourceMgr有这样一个成员:

std::unordered_map<string, Container*> containers_;

container name一般是job name。每个Container里面会放很多从ResourceBase继承来的resource。

  typedef std::pair<uint64, string> ResourceKey;
  typedef std::unordered_map<ResourceKey, ResourceBase*> Container;

单机版无GPU的情况下,默认的device name是"/job:localhost/replica:0/task:0/cpu:0"。其中,job name就是localhost。

单机版的TF会运行每个step的时候,在DirectSession::RunState的构造函数中创建一个per-step的container(ScopedStepContainer),用来存放per-step resources。

ResourceKey里的那个string是resource name,一般是node name。而那个uint64是那个ResourceBase的子类的type的hash_code(type_index.hash_code())。所以不同类型的Resource可以重名。所以,对于某个特定的session来说,它是通过(device name, resource type, container name, resource name) 这样的四元组来定位到一个resource。

DT_RESOURCE是一种新的tensor类型,大约是在去年10月份的时候,0.12版加入的。它所对应的数据类型是ResourceHandle protobuf。该对象的创建以及访问被封装在ResourceMgr类中。其他类不应直接构造此对象。ResourceHandle可以认为是一个per-step的对象。一旦跨step,这个handle可能就会失效。

在DT_RESOURCE出现之前,旧的做法是在OP的内部构造一个PersistentTensor和一个mutex与之对应。这个PersistentTensor的数据类型是string,shape是{2}。里面只放了两样东西container name和shared_name。所以,相比而下,ResourceHandle的信息就要丰富的多了。另外,因为用到了PersistentTensor,所以这种OP必须是stateful。

GetResourceFromContext vs LookupResource:
使用resource datatype的新代码就直接用LookupResource好了。GetResourceFromContext是为了兼容老代码。查询出来的resource对象在用完之后要记得unref。比较干净的做法是这样:

SomeResourceType* res;
OP_REQUIRES_OK(ctx,
                   LookupResource(ctx, HandleFromInput(ctx, 0), &res));
core::ScopedUnref unref_res(res);

我个人认为ResourceOpKernel这个类的历史包袱有点重。新的OP最好不要从这个类继承。它的mu_可以换成call_once,那个PersistentTensor可以被删掉。



此博客中的热门博文

少写代码,多读别人写的代码

在windows下使用llvm+clang

tensorflow distributed runtime初窥