Caffe2 - C++ API
A deep learning, cross platform ML framework
workspace_ops.cc
1 
17 #include "caffe2/core/context.h"
18 #include "caffe2/core/operator.h"
19 
20 namespace caffe2 {
21 namespace {
22 
23 class GetAllBlobNamesOp final : public Operator<CPUContext> {
24  public:
25  GetAllBlobNamesOp(const OperatorDef& operator_def, Workspace* ws)
26  : Operator<CPUContext>(operator_def, ws),
27  include_shared_(GetSingleArgument<int>("include_shared", true)),
28  ws_(ws) {}
29 
30  bool RunOnDevice() override {
31  auto* out = Output(0);
32  const auto& blobs = include_shared_ ? ws_->Blobs() : ws_->LocalBlobs();
33  out->Resize(blobs.size());
34  std::copy(blobs.begin(), blobs.end(), out->mutable_data<std::string>());
35  return true;
36  }
37 
38  private:
39  bool include_shared_;
40  Workspace* ws_;
41 };
42 
43 REGISTER_CPU_OPERATOR(GetAllBlobNames, GetAllBlobNamesOp);
44 OPERATOR_SCHEMA(GetAllBlobNames)
45  .NumInputs(0)
46  .NumOutputs(1)
47  .SetDoc(R"DOC(
48 Return a 1D tensor of strings containing the names
49 of each blob in the active workspace.
50 )DOC")
51  .Arg(
52  "include_shared",
53  "(bool, default true) Whether to include blobs "
54  "inherited from parent workspaces.")
55  .Output(0, "blob_names", "1D tensor of strings containing blob names.");
56 SHOULD_NOT_DO_GRADIENT(GetAllBlobNamesOp);
57 }
58 }
Copyright (c) 2016-present, Facebook, Inc.