tensorflow中有gpu和cpu设备,都可用于计算过程。
目前分析的是一种cpu使用方式,其他的暂时忽略。
使用时每个Executor执行具体任务时device->Compute(op_kernel, &ctx);完成具体计算过程。
这里divice和当前线程的关系,和op_kernel如何完成交互操作关系,待分析。
1、CreateDevices设备完成创建
// TODO(zhifengc/tucker): Figure out the bytes of available RAM.
class ThreadPoolDeviceFactory : public DeviceFactory {
public:
Status ListPhysicalDevices(std::vector* devices) override {
devices->push_back("/physical_device:CPU:0");
return Status::OK();
}
Status CreateDevices(const SessionOptions& options, const string& name_prefix,
std::vector<:unique_ptr>>* devices) override {
int num_numa_nodes = port::NUMANumNodes();
int n = 1;
auto iter = options.config.device_count().find("CPU");
if (iter != options.config.device_count().end()) {
n = iter->second;
}
for (int i = 0; i < n; i++) {
string name = strings::StrCat(name_prefix, "/device:CPU:", i);
std::unique_ptr tpd;
if (options.config.experimental().use_numa_affinity()) {
int numa_node = i % num_numa_nodes;
if (numa_node != i) {
LOG(INFO) << "Only " << num_numa_nodes
<< " NUMA nodes visible in system, "
<< " assigning device " << name << " to NUMA node "
<< numa_node;
}
DeviceLocality dev_locality;
dev_locality.set_numa_node(numa_node);
tpd = absl::make_unique(
options, name, Bytes(256 << 20), dev_locality,
ProcessState::singleton()->GetCPUAllocator(numa_node));
} else {
tpd = absl::make_unique(
options, name, Bytes(256 << 20), DeviceLocality(),
ProcessState::singleton()->GetCPUAllocator(port::kNUMANoAffinity));
}
devices->push_back(std::move(tpd));
}
return Status::OK();
}
};
参考:tensorflow serving 2.2.0代码