Skip to content

SeetaResearch / Dragon

Go to a project
Commit 40e94d24 by Ting PAN
Options

Export Workspace for PyModule

1 parent b35f9320
Expand all
Inline Side-by-side
Showing with 1240 additions and 1621 deletions
CHANGES
------------------------------------------------------------------------ ------------------------------------------------------------------------
The list of most significant changes made over time in Dragon. The list of most significant changes made over time in Dragon.
Dragon 0.3.0.0 (20190309) Dragon 0.3.0.0 (20190402)
DRAGON_VERSION == 3000 DRAGON_VERSION == 3000
Changes (w.r.t. Dragon 0.2.2.13): Changes (w.r.t. Dragon 0.2.2.13):
...@@ -36,6 +36,8 @@ Preview Features: ...@@ -36,6 +36,8 @@ Preview Features:
- The behavior of accumulating gradients have been canceled. - The behavior of accumulating gradients have been canceled.
- Python module now has been assigned to take charge of ``Workspace``.
Bugs fixed: Bugs fixed:
......
Dragon/include/core/graph.h
...@@ -22,17 +22,15 @@ class GraphBase { ...@@ -22,17 +22,15 @@ class GraphBase {
public: public:
/*! \brief Default constructor */ /*! \brief Default constructor */
GraphBase( GraphBase(
const GraphDef& meta_graph, const GraphDef& def,
Workspace* ws); Workspace* ws);
/*! \brief Default deconstructor */ /*! \brief Default deconstructor */
virtual ~GraphBase() {} virtual ~GraphBase() {}
GraphDef BuildUpdateOps(const GraphDef& input_def);
/*! \brief Create a graph from the optimized def */ /*! \brief Create a graph from the optimized def */
virtual bool Create( virtual bool Create(
const GraphDef& optimized_graph, const GraphDef& def,
Workspace* ws) = 0; Workspace* ws) = 0;
/*! \brief Run the graph once synchronously */ /*! \brief Run the graph once synchronously */
...@@ -58,14 +56,14 @@ class GraphBase { ...@@ -58,14 +56,14 @@ class GraphBase {
class Graph : public GraphBase { class Graph : public GraphBase {
public: public:
/*! \brief Default constructor */ /*! \brief Default constructor */
Graph(const GraphDef& meta_graph, Workspace* ws); Graph(const GraphDef& def, Workspace* ws);
/*! \brief Default deconstructor */ /*! \brief Default deconstructor */
virtual ~Graph() { for (auto* op : ops_) delete op; } virtual ~Graph() { for (auto* op : ops_) delete op; }
/*! \brief Create a graph from the optimized def */ /*! \brief Create a graph from the optimized def */
bool Create( bool Create(
const GraphDef& optimized_graph, const GraphDef& def,
Workspace* ws) override; Workspace* ws) override;
/*! \brief Run the graph once synchronously */ /*! \brief Run the graph once synchronously */
......
Dragon/include/core/graph_gradient.h
...@@ -31,7 +31,7 @@ class GraphGradientMaker { ...@@ -31,7 +31,7 @@ class GraphGradientMaker {
const GraphDef& forward_def, const GraphDef& forward_def,
GraphDef& backward_def); GraphDef& backward_def);
void Share(GraphDef& graph); GraphDef Share(const GraphDef& input_def);
void SetTerms(const Map<string, string>& terms) { terms_ = terms; } void SetTerms(const Map<string, string>& terms) { terms_ = terms; }
void SetOperatorPrefix(const string& prefix) { op_prefix_ = prefix; } void SetOperatorPrefix(const string& prefix) { op_prefix_ = prefix; }
......
Dragon/include/core/tensor.h
...@@ -42,9 +42,9 @@ class Tensor { ...@@ -42,9 +42,9 @@ class Tensor {
d = dims[i]; strides_[i] = (int64_t)new_size; d = dims[i]; strides_[i] = (int64_t)new_size;
CHECK_GE(d, 0); CHECK_GE(d, 0);
if (d > 0) new_size *= d; if (d > 0) new_size *= d;
} if (own_mem_) { }
if (size_ != new_size && if (own_mem_) {
capacity_ < new_size * meta_.itemsize()) { if (capacity_ < new_size * meta_.itemsize()) {
memory_.reset(); memory_.reset();
capacity_ = 0; capacity_ = 0;
} }
......
Dragon/include/core/workspace.h
...@@ -29,23 +29,28 @@ class Workspace { ...@@ -29,23 +29,28 @@ class Workspace {
typedef Map<string, unique_ptr<OperatorBase> > OperatorMap; typedef Map<string, unique_ptr<OperatorBase> > OperatorMap;
typedef Map<string, unique_ptr<GraphBase> > GraphMap; typedef Map<string, unique_ptr<GraphBase> > GraphMap;
typedef Map<string, Workspace*> WorkspaceMap;
/*! \brief Constructor */ /*! \brief Constructor */
Workspace(const string& name) : name_(name) { InitWorkspace(); } Workspace(const string& name) : name_(name) { Initialize(); }
/*! \brief Return the name of this workspace */ /*! \brief Return the name of this workspace */
const string& name() { return name_; } const string& name() { return name_; }
/*! \brief Create some internal tensors */ /*! \brief Return the name of stored tensors */
void InitWorkspace(); vector<string> tensors() const;
/*! \brief Return the name of stored graphs */
vector<string> graphs() const;
/*! \brief Move a external workspace into this workspace */ /*! \brief Create some internal tensors */
Workspace* Move(Workspace* ws); void Initialize();
/*! \brief Destory all the tensors */ /*! \brief Destory all the tensors */
void Clear(); void Clear();
/*! \brief Merge from a external workspace */
void MergeFrom(Workspace* ws);
/*! \brief Query the real name of specified tensor */ /*! \brief Query the real name of specified tensor */
string GetTensorName(const string& name) const; string GetTensorName(const string& name) const;
...@@ -66,14 +71,11 @@ class Workspace { ...@@ -66,14 +71,11 @@ class Workspace {
/*! \brief Reset the specified tensor */ /*! \brief Reset the specified tensor */
void ResetTensor(const string& name); void ResetTensor(const string& name);
/*! \brief Return all the stored tensor names */
vector<string> GetTensors() const;
/* \brief Whether the specified filler is in this workspace */ /* \brief Whether the specified filler is in this workspace */
bool HasFiller(const string& name, bool use_remote = true) const; bool HasFiller(const string& name, bool use_remote = true) const;
/*! \brief Create the specified filler */ /*! \brief Create the specified filler */
void CreateFiller(const TensorFillerProto filler); void CreateFiller(const TensorFillerProto& filler);
/*! \brief Return the specified filler */ /*! \brief Return the specified filler */
const TensorFillerProto* GetFiller(const string& name) const; const TensorFillerProto* GetFiller(const string& name) const;
...@@ -82,27 +84,26 @@ class Workspace { ...@@ -82,27 +84,26 @@ class Workspace {
template <class Context> template <class Context>
vector<void*> caches(const vector<size_t>& segments) { vector<void*> caches(const vector<size_t>& segments) {
int64_t nbytes = 0; int64_t nbytes = 0;
vector<void*> ret(segments.size());
for (auto& segment : segments) nbytes += (int64_t)segment; for (auto& segment : segments) nbytes += (int64_t)segment;
Tensor* cache_t = CreateTensor("/share/cache"); auto* T = CreateTensor("/share/cache")->Reshape({ nbytes });
cache_t->Reshape({ nbytes }); ret[0] = T->template mutable_data<uint8_t, Context>();
vector<void*> Bcaches(segments.size());
Bcaches[0] = cache_t->template mutable_data<uint8_t, Context>();
for (int i = 1; i < segments.size(); i++) for (int i = 1; i < segments.size(); i++)
Bcaches[i] = (uint8_t*)Bcaches[i - 1] + segments[i - 1]; ret[i] = (uint8_t*)ret[i - 1] + segments[i - 1];
return Bcaches; return ret;
} }
/*! \brief Create temporal cache segments with the specified type */ /*! \brief Create temporal cache segments with the specified type */
template <typename T, class Context> template <typename T, class Context>
vector<T*> caches(const vector<int64_t>& segments) { vector<T*> caches(const vector<int64_t>& segments) {
vector<size_t> Tsegments; vector<size_t> segments_in_byte;
for (auto& segment : segments) vector<T*> ret(segments.size());
Tsegments.emplace_back(segment * sizeof(T)); for (const auto& e : segments)
vector<void*> Bcaches = caches<Context>(Tsegments); segments_in_byte.emplace_back(e * sizeof(T));
vector<T*> Tcaches(segments.size()); auto ret_in_byte = caches<Context>(segments_in_byte);
for (int i = 0; i < segments.size(); i++) for (int i = 0; i < segments.size(); i++)
Tcaches[i] = (T*)Bcaches[i]; ret[i] = (T*)ret_in_byte[i];
return Tcaches; return ret;
} }
/*! \brief Create a operator in this workspace */ /*! \brief Create a operator in this workspace */
...@@ -124,9 +125,6 @@ class Workspace { ...@@ -124,9 +125,6 @@ class Workspace {
const string& exclude, const string& exclude,
int stream_id = 0); int stream_id = 0);
/*! \brief Return all the stored graph names */
vector<string> GetGraphs() const;
/* \brief Set an alias for the tensor */ /* \brief Set an alias for the tensor */
bool SetTensorAlias(const string& name, const string& alias); bool SetTensorAlias(const string& name, const string& alias);
...@@ -160,7 +158,7 @@ class Workspace { ...@@ -160,7 +158,7 @@ class Workspace {
GraphMap graph_map_; GraphMap graph_map_;
/*! \brief Store the remote workspaces */ /*! \brief Store the remote workspaces */
WorkspaceMap workspace_map_; vector<Workspace*> remote_workspaces_;
}; };
} // namespace dragon } // namespace dragon
......
Dragon/include/operators/misc/gradient_op.h
...@@ -40,8 +40,11 @@ class GradientGatherOp final : public Operator<Context> { ...@@ -40,8 +40,11 @@ class GradientGatherOp final : public Operator<Context> {
public: public:
GradientGatherOp(const OperatorDef& def, Workspace* ws) GradientGatherOp(const OperatorDef& def, Workspace* ws)
: Operator<Context>(def, ws) { : Operator<Context>(def, ws) {
for (int i = 0; i < InputSize(); i++) for (int i = 0; i < InputSize(); i++) {
if (Input(i).name() != "NULL") indices.push_back(i); if (Input(i).name() != "NULL") {
indices.push_back(i);
}
}
} }
USE_OPERATOR_FUNCTIONS; USE_OPERATOR_FUNCTIONS;
...@@ -53,6 +56,16 @@ class GradientGatherOp final : public Operator<Context> { ...@@ -53,6 +56,16 @@ class GradientGatherOp final : public Operator<Context> {
}; };
template <class Context> template <class Context>
class GradientAddOp final : public Operator<Context> {
public:
USE_SIMPLE_CTOR_DTOR(GradientAddOp);
USE_OPERATOR_FUNCTIONS;
void RunOnDevice() override;
template <typename T> void RunWithType();
};
template <class Context>
class StopGradientOp final : public Operator<Context> { class StopGradientOp final : public Operator<Context> {
public: public:
USE_SIMPLE_CTOR_DTOR(StopGradientOp); USE_SIMPLE_CTOR_DTOR(StopGradientOp);
......
Dragon/include/utils/op_kernel.h
...@@ -1033,7 +1033,6 @@ void MixedPrecisionUpdate( ...@@ -1033,7 +1033,6 @@ void MixedPrecisionUpdate(
template <typename T, class Context> template <typename T, class Context>
void BiasAdd( void BiasAdd(
const int count,
const int outer_dim, const int outer_dim,
const int dim, const int dim,
const int inner_dim, const int inner_dim,
......
Dragon/modules/cxx/dragon.cc
...@@ -38,7 +38,7 @@ Workspace* ResetWorkspace(const std::string& name) { ...@@ -38,7 +38,7 @@ Workspace* ResetWorkspace(const std::string& name) {
g_workspaces[name].reset(new Workspace(name)); g_workspaces[name].reset(new Workspace(name));
for (auto& sub_workspace : sub_workspaces[name]) { for (auto& sub_workspace : sub_workspaces[name]) {
if (g_workspaces.count(sub_workspace) > 0) if (g_workspaces.count(sub_workspace) > 0)
g_workspaces[name]->Move( g_workspaces[name]->MergeFrom(
g_workspaces[sub_workspace].get()); g_workspaces[sub_workspace].get());
} }
return g_workspaces[name].get(); return g_workspaces[name].get();
...@@ -55,10 +55,10 @@ void MoveWorkspace( ...@@ -55,10 +55,10 @@ void MoveWorkspace(
std::unique_lock<std::mutex> lock(g_mutex); std::unique_lock<std::mutex> lock(g_mutex);
CHECK(src) << "\nGiven source workspace is invalid."; CHECK(src) << "\nGiven source workspace is invalid.";
CHECK(dst) << "\nGiven destination workspace is invalid."; CHECK(dst) << "\nGiven destination workspace is invalid.";
dst->Move(src); dst->MergeFrom(src);
sub_workspaces[dst->name()].push_back(src->name()); sub_workspaces[dst->name()].push_back(src->name());
LOG(INFO) << "Move the Workspace(" << src->name() << ") " LOG(INFO) << "Move the Workspace(" << src->name() << ") "
<< "into the Workspace(" << dst->name() << ")."; << "into the Workspace(" << dst->name() << ").";
} }
void DestroyWorkspace(const std::string& name) { void DestroyWorkspace(const std::string& name) {
......
Dragon/modules/python/py_autograd.h
...@@ -36,29 +36,6 @@ void AddGradientMethods(pybind11::module& m) { ...@@ -36,29 +36,6 @@ void AddGradientMethods(pybind11::module& m) {
vector<pybind11::bytes>, vector<string>, vector<float> vector<pybind11::bytes>, vector<string>, vector<float>
>(grad_ops, grad.g_inputs, grad.defaults); >(grad_ops, grad.g_inputs, grad.defaults);
}); });
m.def("FlowGradients", [](
const vector<OperatorDef*>& forward_ops,
const vector<string>& targets,
const vector<string>& input_grads,
const vector<string>& ignore_grads,
const bool is_sharing,
const bool verbose) {
// Make => Optimize => Run
GraphDef backward_ops;
GraphGradientMaker maker;
for (auto& grad : input_grads) maker.AddExternalGrad(grad);
for (auto& grad : ignore_grads) maker.AddIgnoreGrad(grad);
maker.Make(forward_ops, targets, backward_ops);
if (is_sharing) maker.Share(backward_ops);
pybind11::gil_scoped_release g;
for (auto& op : backward_ops.op()) {
if (op.type().empty()) continue;
if (verbose) std::cout << op.DebugString() << std::endl;
if (op.has_uid()) ws()->RunOperator(op);
else ws()->RunOperatorOnce(op);
}
});
} }
} // namespace python } // namespace python
......
Dragon/modules/python/py_dragon.h
...@@ -16,15 +16,17 @@ ...@@ -16,15 +16,17 @@
#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
#include "py_types.h" #include "py_types.h"
#include "core/common.h" #include "core/common.h"
#include "core/registry.h"
#include "core/context.h" #include "core/context.h"
#include "core/context_cuda.h"
#include "core/operator.h" #include "core/operator.h"
#include "core/operator_gradient.h" #include "core/registry.h"
#include "core/graph_gradient.h"
#include "core/workspace.h" #include "core/workspace.h"
#include "core/context_cuda.h"
#include "core/graph_gradient.h"
#include "core/operator_gradient.h"
#include "utils/caffemodel.h" #include "utils/caffemodel.h"
#include "onnx/onnx_backend.h"
#include <pybind11/stl.h> #include <pybind11/stl.h>
#include <pybind11/pybind11.h> #include <pybind11/pybind11.h>
...@@ -136,8 +138,6 @@ class NumpyFeeder : public TensorFeederBase { ...@@ -136,8 +138,6 @@ class NumpyFeeder : public TensorFeederBase {
} }
}; };
Workspace* ws();
} // namespace python } // namespace python
} // namespace dragon } // namespace dragon
......
Dragon/modules/python/py_graph.h deleted 100644 → 0
/*!
* Copyright (c) 2017-present, SeetaTech, Co.,Ltd.
*
* Licensed under the BSD 2-Clause License.
* You should have received a copy of the BSD 2-Clause License
* along with the software. If not, See,
*
* <https://opensource.org/licenses/BSD-2-Clause>
*
* ------------------------------------------------------------
*/
#ifndef DRAGON_PYTHON_PY_GRAPH_H_
#define DRAGON_PYTHON_PY_GRAPH_H_
#include "py_dragon.h"
namespace dragon {
namespace python {
void AddGraphMethods(pybind11::module& m) {
/*! \brief Create a graph from the serialized def */
m.def("CreateGraph", [](
const string& serialized,
const bool verbose) {
GraphDef graph_def;
if (!graph_def.ParseFromString(serialized))
LOG(FATAL) << "Failed to parse the GraphDef.";
auto* graph = ws()->CreateGraph(graph_def);
if (verbose) {
// It is not a good design to print the debug string
auto* graph_tensor = ws()->CreateTensor(
"/graph_def/optimized/" + graph->name());
if (graph_tensor->count() > 0) {
auto* data = graph_tensor->mutable_data<string, CPUContext>();
std::cout << data[0] << std::endl;
}
}
// Return the graph name may be different from the def
// We will make a unique dummy name on creating the graph
return graph->name();
});
/*! \brief Run an existing graph */
m.def("RunGraph", [](
const string& name,
const string& include,
const string& exclude) {
pybind11::gil_scoped_release g;
ws()->RunGraph(name, include, exclude);
});
/*! \brief List all of the existing graphs */
m.def("Graphs", []() { ws()->GetGraphs(); });
}
} // namespace python
} // namespace dragon
#endif // DRAGON_PYTHON_PY_GRAPH_H_
\ No newline at end of file
Dragon/modules/python/py_io.h deleted 100644 → 0
/*!
* Copyright (c) 2017-present, SeetaTech, Co.,Ltd.
*
* Licensed under the BSD 2-Clause License.
* You should have received a copy of the BSD 2-Clause License
* along with the software. If not, See,
*
* <https://opensource.org/licenses/BSD-2-Clause>
*
* ------------------------------------------------------------
*/
#ifndef DRAGON_PYTHON_PY_IO_H_
#define DRAGON_PYTHON_PY_IO_H_
#include "py_dragon.h"
namespace dragon {
namespace python {
void AddIOMethods(pybind11::module& m) {
m.def("Snapshot", [](
const string& filename,
vector<string>& names,
const int format) {
vector<Tensor*> tensors;
switch (format) {
case 0: // Pickle
LOG(FATAL) << "Format depends on Pickle. "
"Can't be used in C++.";
break;
case 1: // CaffeModel
for (const auto& e : names)
tensors.emplace_back(ws()->GetTensor(e));
SavaCaffeModel(filename, tensors);
break;
default:
LOG(FATAL) << "Unknwon format, code: " << format;
}
});
m.def("Restore", [](
const string& filename,
const int format) {
switch (format) {
case 0: // Pickle
LOG(FATAL) << "Format depends on Pickle. "
"Can't be used in C++.";
break;
case 1: // CaffeModel
LoadCaffeModel(filename, ws());
break;
default:
LOG(FATAL) << "Unknwon format, code: " << format;
}
});
}
} // namespace python
} // namespace dragon
#endif // DRAGON_PYTHON_PY_IO_H_
\ No newline at end of file
Dragon/modules/python/py_onnx.h deleted 100644 → 0
/*!
* Copyright (c) 2017-present, SeetaTech, Co.,Ltd.
*
* Licensed under the BSD 2-Clause License.
* You should have received a copy of the BSD 2-Clause License
* along with the software. If not, See,
*
* <https://opensource.org/licenses/BSD-2-Clause>
*
* ------------------------------------------------------------
*/
#ifndef DRAGON_PYTHON_PY_ONNX_H_
#define DRAGON_PYTHON_PY_ONNX_H_
#include "onnx/onnx_backend.h"
#include "py_dragon.h"
namespace dragon {
namespace python {
void AddONNXMethods(pybind11::module& m) {
m.def("ImportONNXModel", [](
const string& model_path) {
GraphDef init_graph, pred_graph;
onnx::ONNXBackend onnx_backend;
onnx_backend.Prepare(model_path, &init_graph, &pred_graph);
// Serializing to Python is intractable
// We should apply the initializer immediately
ws()->CreateGraph(init_graph);
ws()->RunGraph(init_graph.name(), "", "");
return pybind11::bytes(pred_graph.SerializeAsString());
});
}
} // namespace python
} // namespace dragon
#endif // DRAGON_PYTHON_PY_ONNX_H_
\ No newline at end of file
Dragon/modules/python/py_operator.h
...@@ -20,36 +20,14 @@ namespace dragon { ...@@ -20,36 +20,14 @@ namespace dragon {
namespace python { namespace python {
void AddOperatorMethods(pybind11::module& m) { void AddOperatorMethods(pybind11::module& m) {
/*! \brief Return all the registered operators */ /*! \brief Return the registered operators */
m.def("RegisteredOperators", []() { return CPUOperatorRegistry()->keys(); }); m.def("RegisteredOperators", []() {
return CPUOperatorRegistry()->keys();
/*! \brief Return all the operators without gradients */
m.def("NoGradientOperators", []() { return NoGradientRegistry()->keys(); });
/*! \brief Run a operator from the def reference */
m.def("RunOperator", [](
OperatorDef* def,
const bool verbose) {
pybind11::gil_scoped_release g;
if (verbose) {
// It is not a good design to print the debug string
std::cout << def->DebugString() << std::endl;
}
ws()->RunOperator(*def);
}); });
/*! \brief Run a operator from the serialized def */ /*! \brief Return the non-gradient operators */
m.def("RunOperator", []( m.def("NoGradientOperators", []() {
const string& serialized, return NoGradientRegistry()->keys();
const bool verbose) {
OperatorDef def;
CHECK(def.ParseFromString(serialized));
pybind11::gil_scoped_release g;
if (verbose) {
// It is not a good design to print the debug string
std::cout << def.DebugString() << std::endl;
}
ws()->RunOperatorOnce(def);
}); });
} }
......
Dragon/modules/python/py_tensor.h
...@@ -22,208 +22,51 @@ namespace python { ...@@ -22,208 +22,51 @@ namespace python {
void AddTensorMethods(pybind11::module& m) { void AddTensorMethods(pybind11::module& m) {
/*! \brief Export the Tensor class */ /*! \brief Export the Tensor class */
pybind11::class_<Tensor>(m, "Tensor") pybind11::class_<Tensor>(m, "Tensor")
/*! \brief Return the number of dimensions */
.def_property_readonly("ndim", &Tensor::ndim) .def_property_readonly("ndim", &Tensor::ndim)
/*! \brief Return all the dimensions */
.def_property_readonly("dims", &Tensor::dims) .def_property_readonly("dims", &Tensor::dims)
/*! \brief Return the total number of elements */
.def_property_readonly("size", &Tensor::size) .def_property_readonly("size", &Tensor::size)
/*! \brief Return the data type */
.def_property_readonly("dtype", [](Tensor* self) { .def_property_readonly("dtype", [](Tensor* self) {
return TypeMetaToString(self->meta()); return TypeMetaToString(self->meta());
}).def_property_readonly("device", [](Tensor* self) { })
/*! \brief Return the device information */
.def_property_readonly("device", [](Tensor* self) {
if (self->has_memory()) { if (self->has_memory()) {
Map<string, string> mem_info = self->memory()->info(); auto mem_info = self->memory()->info();
return std::tuple<string, int>( return std::tuple<string, int>(
mem_info["device_type"], atoi( mem_info["device_type"], atoi(
mem_info["device_id"].c_str())); mem_info["device_id"].c_str()));
} else { } else {
return std::tuple<string, int>("Unknown", 0); return std::tuple<string, int>("Unknown", 0);
} }
}).def("ToCPU", [](Tensor* self) { })
CHECK(self->has_memory()) << "\nTensor(" << self->name()
<< ") does not initialize or had been reset."; /*! \brief Switch the memory to the cpu context */
.def("ToCPU", [](Tensor* self) {
CHECK(self->has_memory())
<< "\nTensor(" << self->name() << ") "
<< "does not initialize or had been reset.";
self->memory()->ToCPU(); self->memory()->ToCPU();
}).def("ToCUDA", [](Tensor* self, const int device_id) { })
/*! \brief Switch the memory to the cuda context */
.def("ToCUDA", [](Tensor* self, int device_id) {
#ifdef WITH_CUDA #ifdef WITH_CUDA
CHECK(self->has_memory()) << "\nTensor(" << self->name() CHECK(self->has_memory())
<< ") does not initialize or had been reset."; << "\nTensor(" << self->name() << ") "
<< "does not initialize or had been reset.";
self->memory()->SwitchToCUDADevice(device_id); self->memory()->SwitchToCUDADevice(device_id);
#else #else
CUDA_NOT_COMPILED; CUDA_NOT_COMPILED;
#endif #endif
}); });
/*! \brief List all the existing tensors */
m.def("Tensors", []() { return ws()->GetTensors(); });
/*! \brief Indicate whether the given tensor is existing */
m.def("HasTensor", [](
const string& name) -> bool {
return ws()->HasTensor(name);
});
/*! \brief Return the unique name of given tensor */
m.def("GetTensorName", [](
const string& name) -> string {
return ws()->GetTensorName(name);
});
/*! \brief Create a tensor with the given name */
m.def("CreateTensor", [](
const string& name) -> void {
ws()->CreateTensor(name);
});
/*! \brief Create a tensor with the given name */
m.def("ResetTensor", [](
const string& name) -> void {
ws()->ResetTensor(name);
});
/*! \brief Create a tensor with the given shape */
m.def("TensorFromShape", [](
const string& name,
const vector<int64_t>& shape,
const string& dtype) {
const TypeMeta& meta = TypeStringToMeta(dtype);
if (meta.id() == 0) {
LOG(FATAL) << "Unsupported data type: " + dtype + ".";
}
Tensor* tensor = ws()->CreateTensor(name);
if (meta.id() != tensor->meta().id() && tensor->meta().id() != 0)
LOG(WARNING) << "Set Tensor(" << tensor->name() << ")"
<< " with different data type from original one.";
tensor->Reshape(shape);
tensor->raw_mutable_data<CPUContext>(meta);
});
/*! \brief Create a tensor with the given array */
m.def("TensorFromPyArray", [](
const string& name,
pybind11::object py_array) {
PyArrayObject* array = PyArray_GETCONTIGUOUS(
reinterpret_cast<PyArrayObject*>(py_array.ptr()));
const TypeMeta& meta = TypeNPYToMeta(PyArray_TYPE(array));
if (meta.id() == 0) LOG(FATAL) << "Unsupported data type.";
Tensor* tensor = ws()->CreateTensor(name);
tensor->SetMeta(meta);
int ndim = PyArray_NDIM(array);
npy_intp* npy_dims = PyArray_DIMS(array);
vector<int64_t> dims;
for (int i = 0; i < ndim; i++) dims.push_back(npy_dims[i]);
tensor->Reshape(dims);
auto* data = static_cast<void*>(PyArray_DATA(array));
if (!tensor->has_memory()) {
MixedMemory* memory(new MixedMemory());
memory->set_cpu_data(data, tensor->nbytes());
tensor->set_memory(memory);
} else {
if (tensor->DECREFPyArray) tensor->DECREFPyArray();
tensor->memory()->set_cpu_data(data, tensor->nbytes());
}
// Follow the codes of PyTorch
// Here we bind the DECREF to Tensor
// ResetTensor() or ResetWorkspace() can trigger it
tensor->DECREFPyArray = [array]()->void { Py_XDECREF(array); };
});
/*! \brief Create a tensor copied from an existing one */
m.def("TensorFromTensor", [](
const string& name,
const string& other,
const string& dev1,
const string& dev2) {
DeviceOption dst_ctx, src_ctx;
dst_ctx.ParseFromString(dev1);
src_ctx.ParseFromString(dev2);
Tensor* srcT = ws()->GetTensor(other);
Tensor* dstT = ws()->CreateTensor(name);
dstT->ReshapeLike(*srcT);
const TypeMeta& meta = srcT->meta();
if (dst_ctx.device_type() == PROTO_CUDA) {
if (src_ctx.device_type() == PROTO_CUDA) {
// CUDA <- CUDA
CUDAContext::MemcpyEx<CUDAContext, CUDAContext>(
srcT->nbytes(),
dstT->raw_mutable_data<CUDAContext>(meta),
srcT->raw_data<CUDAContext>(),
src_ctx.device_id());
} else {
// CUDA <- CPU
CUDAContext::MemcpyEx<CUDAContext, CPUContext>(
srcT->nbytes(),
dstT->raw_mutable_data<CUDAContext>(meta),
srcT->raw_data<CPUContext>(),
dst_ctx.device_id());
}
} else {
if (src_ctx.device_type() == PROTO_CUDA) {
// CPU <- CUDA
CUDAContext::MemcpyEx<CPUContext, CUDAContext>(
srcT->nbytes(),
dstT->raw_mutable_data<CPUContext>(meta),
srcT->raw_data<CUDAContext>(),
src_ctx.device_id());
} else {
// CPU <- CPU
CPUContext::Memcpy<CUDAContext, CUDAContext>(
srcT->nbytes(),
dstT->raw_mutable_data<CPUContext>(meta),
srcT->raw_data<CPUContext>());
}
}
});
/*! \brief Return a array zero-copied from an existing tensor */
m.def("TensorToPyArray", [](
const string& name,
const bool readonly) {
Tensor* tensor = ws()->GetTensor(name);
CHECK_GT(tensor->count(), 0);
vector<npy_intp> dims;
for (const auto dim : tensor->dims()) dims.push_back(dim);
int npy_type = TypeMetaToNPY(tensor->meta());
if (npy_type == -1) {
LOG(FATAL) << "Tensor(" + tensor->name() + ") "
"with dtype." + TypeMetaToString(tensor->meta()) +
" is not supported by numpy.";
}
auto* data = readonly ?
const_cast<void*>(tensor->raw_data<CPUContext>()) :
tensor->raw_mutable_data<CPUContext>();
PyObject* array = PyArray_SimpleNewFromData(
tensor->ndim(), dims.data(), npy_type, data);
return pybind11::reinterpret_steal<pybind11::object>(array);
});
/*! \brief Create a tensor from the specified filler */
m.def("CreateFiller", [](
const string& serialized) {
TensorFillerProto filler_proto;
if (!filler_proto.ParseFromString(serialized))
LOG(FATAL) << "Failed to parse the TensorFiller.";
ws()->CreateFiller(filler_proto);
ws()->CreateTensor(filler_proto.tensor());
});
/*! \brief Return the filler type of a tensor */
m.def("GetFillerType", [](const string& name) {
return ws()->GetFiller(name)->type();
});
/* \brief Set an alias for the tensor */
m.def("SetTensorAlias", [](
const string& name,
const string& alias) {
if (!ws()->HasTensor(name)) {
LOG(FATAL) << "Tensor(" + name << ") has not "
"been registered in the current workspace.";
}
ws()->SetTensorAlias(name, alias);
});
/*! \brief Return the CXX Tensor reference */
m.def("GetTensor", [](
const string& name) {
return ws()->GetTensor(name);
}, pybind11::return_value_policy::reference_internal);
} }
} // namespace python } // namespace python
......
Dragon/python/dragon/__init__.py
...@@ -22,6 +22,9 @@ import dragon.config as config ...@@ -22,6 +22,9 @@ import dragon.config as config
# Core # Core
from dragon.core.tensor import Tensor from dragon.core.tensor import Tensor
import dragon.core.workspace as workspace import dragon.core.workspace as workspace
from dragon.core.workspace import Workspace
from dragon.core.workspace import get_default_workspace
from dragon.core.workspace import reset_default_workspace
import dragon.core.tensor_utils as tensor_utils import dragon.core.tensor_utils as tensor_utils
import dragon.core.mpi as mpi import dragon.core.mpi as mpi
import dragon.core.cuda as cuda import dragon.core.cuda as cuda
...@@ -41,7 +44,6 @@ from dragon.vm.theano.tensor import grad as grad ...@@ -41,7 +44,6 @@ from dragon.vm.theano.tensor import grad as grad
from dragon.core.scope import name_scope, get_default_name_scope from dragon.core.scope import name_scope, get_default_name_scope
from dragon.core.scope import phase_scope, get_default_phase from dragon.core.scope import phase_scope, get_default_phase
from dragon.core.scope import device_scope, get_default_device from dragon.core.scope import device_scope, get_default_device
from dragon.core.scope import WorkspaceScope as ws_scope
# Version # Version
from dragon.version import version from dragon.version import version
......
Dragon/python/dragon/config.py
...@@ -15,8 +15,9 @@ from __future__ import absolute_import ...@@ -15,8 +15,9 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import dragon.import_c_api as C from dragon import import_c_api as _C
import dragon.core.logging as logging from dragon.core import logging as _logging
option = {} option = {}
...@@ -290,12 +291,12 @@ def SetLoggingLevel(level): ...@@ -290,12 +291,12 @@ def SetLoggingLevel(level):
The default level is *INFO*. The default level is *INFO*.
""" """
C.SetLoggingLevel(level) _C.SetLoggingLevel(level)
logging.set_verbosity({ _logging.set_verbosity({
'DEBUG': logging.DEBUG, 'DEBUG': _logging.DEBUG,
'INFO': logging.INFO, 'INFO': _logging.INFO,
'WARNING': logging.WARN, 'WARNING': _logging.WARN,
'ERROR': logging.ERROR, 'ERROR': _logging.ERROR,
'FATAL': logging.FATAL, 'FATAL': _logging.FATAL,
}[level] }[level]
) )
\ No newline at end of file
Dragon/python/dragon/core/cuda.py
...@@ -15,7 +15,7 @@ from __future__ import absolute_import ...@@ -15,7 +15,7 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import dragon.import_c_api as _C from dragon import import_c_api as _C
def IsCUDADriverSufficient(): def IsCUDADriverSufficient():
......
Dragon/python/dragon/core/gradient_maker.py
...@@ -30,11 +30,10 @@ from __future__ import print_function ...@@ -30,11 +30,10 @@ from __future__ import print_function
from collections import defaultdict from collections import defaultdict
import dragon.proto.dragon_pb2 as pb from dragon import import_c_api as _C
import dragon.import_c_api as C from dragon.core import helper as _helper
from dragon.proto import dragon_pb2 as _proto_def
from dragon.core.helper import OperatorHelper from dragon.core import proto_utils as _proto_utils
from dragon.core.proto_utils import MakeOperatorDef
class GraphGradientMaker(object): class GraphGradientMaker(object):
...@@ -62,16 +61,22 @@ class GraphGradientMaker(object): ...@@ -62,16 +61,22 @@ class GraphGradientMaker(object):
The OpDef, outputs and defaults of ``BackwardOp``. The OpDef, outputs and defaults of ``BackwardOp``.
""" """
g_ops, g_inputs, defaults = C.CreateGradientDefs( g_ops, g_inputs, defaults = _C.CreateGradientDefs(
forward_op.SerializeToString(), g_outputs) forward_op.SerializeToString(), g_outputs)
for idx, g_op in enumerate(g_ops): for idx, g_op in enumerate(g_ops):
new_def = pb.OperatorDef() new_def = _proto_def.OperatorDef()
new_def.ParseFromString(g_op) new_def.ParseFromString(g_op)
g_ops[idx] = new_def g_ops[idx] = new_def
return g_ops, g_inputs, defaults return g_ops, g_inputs, defaults
@classmethod @classmethod
def CheckGrad(cls, forward_op, inputs_to_grads, blacklist, targets): def CheckGrad(
cls,
forward_op,
inputs_to_grads,
blacklist,
targets,
):
"""Check if missing Grads. If True, skip this Op. """Check if missing Grads. If True, skip this Op.
Parameters Parameters
...@@ -91,7 +96,7 @@ class GraphGradientMaker(object): ...@@ -91,7 +96,7 @@ class GraphGradientMaker(object):
The result of checking and generated filling grads. The result of checking and generated filling grads.
""" """
if forward_op.type in C.NO_GRADIENT_OPERATORS: if forward_op.type in _C.NO_GRADIENT_OPERATORS:
for input in forward_op.input: blacklist.add(input) for input in forward_op.input: blacklist.add(input)
return True, None return True, None
...@@ -114,7 +119,13 @@ class GraphGradientMaker(object): ...@@ -114,7 +119,13 @@ class GraphGradientMaker(object):
return False, gen_grads return False, gen_grads
@classmethod @classmethod
def Make(cls, forward_ops, targets, input_grads=None, auto_names=True): def Make(
cls,
forward_ops,
targets,
input_grads=None,
auto_names=True,
):
"""Make ``BackwardOps`` based on ``ForwardOps``. """Make ``BackwardOps`` based on ``ForwardOps``.
Parameters Parameters
...@@ -149,7 +160,7 @@ class GraphGradientMaker(object): ...@@ -149,7 +160,7 @@ class GraphGradientMaker(object):
# PLAY for the forward # PLAY for the forward
for forward_op in forward_ops: for forward_op in forward_ops:
if forward_op.type in C.NO_GRADIENT_OPERATORS: continue if forward_op.type in _C.NO_GRADIENT_OPERATORS: continue
outputs = [o for o in forward_op.output] outputs = [o for o in forward_op.output]
for input in forward_op.input: for input in forward_op.input:
if input not in outputs: if input not in outputs:
...@@ -176,14 +187,17 @@ class GraphGradientMaker(object): ...@@ -176,14 +187,17 @@ class GraphGradientMaker(object):
op_inputs.append(item[0]) op_inputs.append(item[0])
op_outputs.append(item[0] + '_grad') op_outputs.append(item[0] + '_grad')
values.append(defaults[item[1]]) values.append(defaults[item[1]])
gen_op = MakeOperatorDef('GradientGenerate', op_inputs, op_outputs, defaults=values) gen_op = _proto_utils.MakeOperatorDef(
gen_op.name = OperatorHelper.get_name() if auto_names else 'runtime' 'GradientGenerate', op_inputs, op_outputs, defaults=values)
gen_op.name = _helper.OperatorHelper. \
get_name() if auto_names else 'runtime'
if forward_op.HasField('device_option'): if forward_op.HasField('device_option'):
gen_op.device_option.CopyFrom(forward_op.device_option) gen_op.device_option.CopyFrom(forward_op.device_option)
backward_ops.append(gen_op) backward_ops.append(gen_op)
# GradientOp # GradientOp
for g_op in g_ops: for g_op in g_ops:
g_op.name = OperatorHelper.get_name() if auto_names else 'runtime' g_op.name = _helper.OperatorHelper. \
get_name() if auto_names else 'runtime'
backward_ops.append(g_op) backward_ops.append(g_op)
# Split & Gather grads for multi-used input # Split & Gather grads for multi-used input
...@@ -208,10 +222,12 @@ class GraphGradientMaker(object): ...@@ -208,10 +222,12 @@ class GraphGradientMaker(object):
for idx in range(grads_count[g_output]): for idx in range(grads_count[g_output]):
if '%s_autosplit_%d' % (g_output, idx) in all_split_grads: if '%s_autosplit_%d' % (g_output, idx) in all_split_grads:
split_inputs.append('%s_autosplit_%d' % (g_output, idx)) split_inputs.append('%s_autosplit_%d' % (g_output, idx))
gather_op = MakeOperatorDef('GradientGather', split_inputs, [g_output]) gather_op = _proto_utils.MakeOperatorDef(
'GradientGather', split_inputs, [g_output])
if g_op.HasField('device_option'): if g_op.HasField('device_option'):
gather_op.device_option.CopyFrom(g_op.device_option) gather_op.device_option.CopyFrom(g_op.device_option)
gather_op.name = OperatorHelper.get_name() if auto_names else 'runtime' gather_op.name = _helper.OperatorHelper. \
get_name() if auto_names else 'runtime'
backward_ops.append(gather_op) backward_ops.append(gather_op)
g_op.output[g_output_idx] = split_name g_op.output[g_output_idx] = split_name
......
Dragon/python/dragon/core/helper.py
...@@ -17,7 +17,8 @@ from __future__ import print_function ...@@ -17,7 +17,8 @@ from __future__ import print_function
import math import math
import numpy import numpy
import dragon
from dragon.core import workspace as _workspace
class OperatorHelper(object): class OperatorHelper(object):
...@@ -39,11 +40,11 @@ class OperatorHelper(object): ...@@ -39,11 +40,11 @@ class OperatorHelper(object):
@classmethod @classmethod
def get_index_and_name(cls, prefix='Op'): def get_index_and_name(cls, prefix='Op'):
name = dragon.workspace.GetDummyName(prefix, domain='Operator') name = _workspace.GetDummyName(prefix, domain='Operator')
try: try:
_, op_idx = name.split('_') _, op_idx = name.split('_')
except: except:
name = dragon.workspace.GetDummyName(prefix, domain='Operator') name = _workspace.GetDummyName(prefix, domain='Operator')
_, op_idx = name.split('_') _, op_idx = name.split('_')
return int(op_idx), name return int(op_idx), name
......
Dragon/python/dragon/core/mpi.py
...@@ -15,7 +15,7 @@ from __future__ import absolute_import ...@@ -15,7 +15,7 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import dragon.import_c_api as _C from dragon import import_c_api as _C
_GLOBAL_MPI_IS_INIT = False _GLOBAL_MPI_IS_INIT = False
......
Dragon/python/dragon/core/proto_utils.py
...@@ -9,7 +9,7 @@ ...@@ -9,7 +9,7 @@
# #
# ------------------------------------------------------------ # ------------------------------------------------------------
"""Define some helpful protobuf makers here.""" """Define some helpful protocol buffer makers here."""
from __future__ import absolute_import from __future__ import absolute_import
from __future__ import division from __future__ import division
...@@ -17,28 +17,28 @@ from __future__ import print_function ...@@ -17,28 +17,28 @@ from __future__ import print_function
import sys import sys
import copy import copy
import numpy as np import numpy
from google.protobuf.message import Message
import dragon.config as cfg from dragon import config as _cfg
import dragon.import_c_api as _C from dragon import import_c_api as _C
from dragon.proto import dragon_pb2 as pb from dragon.core import scope as _scope
from dragon.core.scope import get_default_device from dragon.proto import dragon_pb2 as _proto_def
from google.protobuf.message import Message as _Message
if sys.version_info >= (3,0): if sys.version_info >= (3,0):
def MakeArgument(key, value): def MakeArgument(key, value):
argument = pb.Argument() argument = _proto_def.Argument()
argument.name = key argument.name = key
if type(value) is float: argument.f = value if type(value) is float: argument.f = value
elif type(value) in (bool, int, np.int64) : argument.i = value elif type(value) in (bool, int, numpy.int64) : argument.i = value
elif type(value) is bytes: argument.s = value elif type(value) is bytes: argument.s = value
elif type(value) is str: argument.s = str.encode(value) elif type(value) is str: argument.s = str.encode(value)
elif isinstance(value, Message): argument.s = value.SerializeToString() elif isinstance(value, _Message): argument.s = value.SerializeToString()
elif all(type(v) is float for v in value): argument.floats.extend(value) elif all(type(v) is float for v in value): argument.floats.extend(value)
elif all(type(v) is int for v in value): argument.ints.extend(value) elif all(type(v) is int for v in value): argument.ints.extend(value)
elif all(type(v) is str for v in value): argument.strings.extend([str.encode(v) for v in value]) elif all(type(v) is str for v in value): argument.strings.extend([str.encode(v) for v in value])
elif all(isinstance(v, Message) for v in value): elif all(isinstance(v, _Message) for v in value):
argument.strings.extend([v.SerializeToString() for v in value]) argument.strings.extend([v.SerializeToString() for v in value])
else: else:
raise ValueError( raise ValueError(
...@@ -47,20 +47,20 @@ if sys.version_info >= (3,0): ...@@ -47,20 +47,20 @@ if sys.version_info >= (3,0):
return argument return argument
else: else:
def MakeArgument(key, value): def MakeArgument(key, value):
argument = pb.Argument() argument = _proto_def.Argument()
argument.name = key argument.name = key
if type(value) is float: argument.f = value if type(value) is float: argument.f = value
elif type(value) in (bool, int, long, np.int64) : argument.i = value elif type(value) in (bool, int, long, numpy.int64) : argument.i = value
elif type(value) is str: argument.s = value elif type(value) is str: argument.s = value
elif type(value) is unicode: argument.s = str(value) elif type(value) is unicode: argument.s = str(value)
elif isinstance(value, Message): argument.s = value.SerializeToString() elif isinstance(value, _Message): argument.s = value.SerializeToString()
elif all(type(v) is float for v in value): argument.floats.extend(value) elif all(type(v) is float for v in value): argument.floats.extend(value)
elif all(type(v) is int for v in value): argument.ints.extend(value) elif all(type(v) is int for v in value): argument.ints.extend(value)
elif all(type(v) is long for v in value): argument.ints.extend(value) elif all(type(v) is long for v in value): argument.ints.extend(value)
elif all(type(v) is str for v in value): argument.strings.extend(value) elif all(type(v) is str for v in value): argument.strings.extend(value)
elif all(type(v) is unicode for v in value): elif all(type(v) is unicode for v in value):
argument.strings.extend([str(v) for v in value]) argument.strings.extend([str(v) for v in value])
elif all(isinstance(v, Message) for v in value): elif all(isinstance(v, _Message) for v in value):
argument.strings.extend([v.SerializeToString() for v in value]) argument.strings.extend([v.SerializeToString() for v in value])
else: else:
raise ValueError( raise ValueError(
...@@ -70,10 +70,16 @@ else: ...@@ -70,10 +70,16 @@ else:
def MakeOperatorDef( def MakeOperatorDef(
op_type, inputs=(), outputs=(), op_type,
name='', uid=None, device_option=None, inputs=(),
arg=None, **kwargs): outputs=(),
operator = pb.OperatorDef() name='',
uid=None,
device_option=None,
arg=None,
**kwargs
):
operator = _proto_def.OperatorDef()
operator.type = op_type operator.type = op_type
operator.name = name operator.name = name
operator.input.extend([str(tensor) for tensor in inputs]) operator.input.extend([str(tensor) for tensor in inputs])
...@@ -92,9 +98,15 @@ def MakeOperatorDef( ...@@ -92,9 +98,15 @@ def MakeOperatorDef(
def MakeCXXOperatorDef( def MakeCXXOperatorDef(
op_type, inputs=(), outputs=(), op_type,
name='', uid=None, device_option=None, inputs=(),
arg=None, **kwargs): outputs=(),
name='',
uid=None,
device_option=None,
arg=None,
**kwargs
):
c_def = _C.OperatorDef() c_def = _C.OperatorDef()
py_def = MakeOperatorDef( py_def = MakeOperatorDef(
op_type, inputs, outputs, name, uid, op_type, inputs, outputs, name, uid,
...@@ -104,7 +116,7 @@ def MakeCXXOperatorDef( ...@@ -104,7 +116,7 @@ def MakeCXXOperatorDef(
def MakeDeviceOption(device_type, device_id, rng_seed=None): def MakeDeviceOption(device_type, device_id, rng_seed=None):
option = pb.DeviceOption() option = _proto_def.DeviceOption()
option.device_type = device_type option.device_type = device_type
option.device_id = device_id option.device_id = device_id
if rng_seed is not None: option.random_seed = rng_seed if rng_seed is not None: option.random_seed = rng_seed
...@@ -133,7 +145,7 @@ def GetDeviceOption(device_type, device_id=0, rng_seed=None): ...@@ -133,7 +145,7 @@ def GetDeviceOption(device_type, device_id=0, rng_seed=None):
def GetDefaultDeviceOption(): def GetDefaultDeviceOption():
device_info = get_default_device() device_info = _scope.get_default_device()
if device_info is not None: if device_info is not None:
return GetDeviceOption( return GetDeviceOption(
device_info['device_type'], device_info['device_type'],
...@@ -142,10 +154,10 @@ def GetDefaultDeviceOption(): ...@@ -142,10 +154,10 @@ def GetDefaultDeviceOption():
def GetGlobalDeviceOption(): def GetGlobalDeviceOption():
option = cfg.GetGlobalOptions() options = _cfg.GetGlobalOptions()
return GetDeviceOption( return GetDeviceOption(
option['device'], options['device'],
option['device_id']) options['device_id'])
# Fix the python stdout # Fix the python stdout
...@@ -159,6 +171,5 @@ class Unbuffered(object): ...@@ -159,6 +171,5 @@ class Unbuffered(object):
return getattr(self.stream, attr) return getattr(self.stream, attr)
# Clear the stdout buffer for mpi(C++ && Python) # Clear the stdout buffer for mpi
import sys
sys.stdout = Unbuffered(sys.stdout) sys.stdout = Unbuffered(sys.stdout)
\ No newline at end of file
Dragon/python/dragon/core/scope.py
...@@ -13,92 +13,7 @@ from __future__ import absolute_import ...@@ -13,92 +13,7 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import threading from dragon.core import tls as _tls
import dragon.import_c_api as _C
from contextlib import contextmanager
__all__ = [
'name_scope',
'phase_scope',
'device_scope',
'get_default_phase',
'get_default_device',
'get_default_name_scope',
'WorkspaceScope',
]
class _ThreadLocalStack(threading.local):
def __init__(self):
super(_ThreadLocalStack, self).__init__()
self._enforce_nesting = True
self.stack = []
def get_default(self):
return self.stack[-1] if len(self.stack) >= 1 else None
def is_cleared(self):
return not self.stack
@property
def enforce_nesting(self):
return self._enforce_nesting
@enforce_nesting.setter
def enforce_nesting(self, value):
self._enforce_nesting = value
@contextmanager
def get_controller(self, default):
"""A context manager for manipulating a default stack."""
self.stack.append(default)
try:
yield default
finally:
# stack may be empty if reset() was called
if self.stack:
if self._enforce_nesting:
if self.stack[-1] is not default:
raise AssertionError(
"Nesting violated for default stack of %s objects" %
type(default))
self.stack.pop()
else:
self.stack.remove(default)
class WorkspaceScope(object):
"""WorkspaceScope is a auxiliary to assign the specific workspace.
Examples
--------
>>> import dragon as dg
>>> with WorkspaceScope('session1'): pass
>>> with dg.ws_scope('session2'): pass
"""
def __init__(self, ws_name):
assert isinstance(ws_name, type('str')), \
'WorkspaceScope takes in a string as its argument.'
assert ws_name != '', \
'The workspace name should not be empty.'
self.ws = ws_name
self.prev = 'default'
def __enter__(self):
self.prev = _C.CurrentWorkspace()
_C.SwitchWorkspace(self.ws, True)
def __exit__(self, type, value, traceback):
_C.SwitchWorkspace(self.prev, True)
_GLOBAL_TENSOR_STACK = _ThreadLocalStack()
_GLOBAL_PHASE_STACK = _ThreadLocalStack()
_GLOBAL_DEVICE_STACK = _ThreadLocalStack()
_PREDEFINED_SCOPE_SEPARATOR = '/'
def name_scope(name): def name_scope(name):
...@@ -140,7 +55,7 @@ def device_scope(device_type, device_id=0): ...@@ -140,7 +55,7 @@ def device_scope(device_type, device_id=0):
""" """
device_type, device_id, device_type.lower(), device_id device_type, device_id, device_type.lower(), device_id
assert device_type in ['cpu', 'gpu', 'cuda', 'cnml'] assert device_type in ('cpu', 'gpu', 'cuda', 'cnml')
# Default names # Default names
if device_type == 'gpu': device_type = 'cuda' if device_type == 'gpu': device_type = 'cuda'
return _GLOBAL_DEVICE_STACK.get_controller({ return _GLOBAL_DEVICE_STACK.get_controller({
...@@ -212,4 +127,10 @@ def get_default_device(): ...@@ -212,4 +127,10 @@ def get_default_device():
The device dict. The device dict.
""" """
return _GLOBAL_DEVICE_STACK.get_default() return _GLOBAL_DEVICE_STACK.get_default()
\ No newline at end of file
_GLOBAL_TENSOR_STACK = _tls.Stack()
_GLOBAL_PHASE_STACK = _tls.Stack()
_GLOBAL_DEVICE_STACK = _tls.Stack()
_PREDEFINED_SCOPE_SEPARATOR = '/'
\ No newline at end of file
Dragon/python/dragon/core/tensor_utils.py
...@@ -16,10 +16,10 @@ from __future__ import division ...@@ -16,10 +16,10 @@ from __future__ import division
from __future__ import print_function from __future__ import print_function
import numpy import numpy
import dragon
from dragon.core.tensor import Tensor from dragon.core import workspace as _workspace
from dragon.core.proto_utils import GetDeviceOption from dragon.core import proto_utils as _proto_utils
from dragon.core.tensor import Tensor as _Tensor
def FromShape(shape, dtype='float32', name=None): def FromShape(shape, dtype='float32', name=None):
...@@ -47,9 +47,8 @@ def FromShape(shape, dtype='float32', name=None): ...@@ -47,9 +47,8 @@ def FromShape(shape, dtype='float32', name=None):
tensor.shape = list(shape) tensor.shape = list(shape)
if not isinstance(shape, (tuple, list)): if not isinstance(shape, (tuple, list)):
raise TypeError('The shape should be a tuple or list.') raise TypeError('The shape should be a tuple or list.')
dragon.C.TensorFromShape( _get_workspace().TensorFromShape(
_stringify_tensor(tensor), _stringify_tensor(tensor), list(shape), dtype)
list(shape), dtype)
return tensor return tensor
...@@ -70,7 +69,8 @@ def SetShape(tensor, shape, dtype='float32'): ...@@ -70,7 +69,8 @@ def SetShape(tensor, shape, dtype='float32'):
None None
""" """
dragon.C.TensorFromShape(_stringify_tensor(tensor), shape, dtype) _get_workspace().TensorFromShape(
_stringify_tensor(tensor), shape, dtype)
def FromTensor(src, src_ctx=None, name=None, ctx=None): def FromTensor(src, src_ctx=None, name=None, ctx=None):
...@@ -97,15 +97,17 @@ def FromTensor(src, src_ctx=None, name=None, ctx=None): ...@@ -97,15 +97,17 @@ def FromTensor(src, src_ctx=None, name=None, ctx=None):
""" """
tensor = _try_get_tensor(name) tensor = _try_get_tensor(name)
if src_ctx is None: src_ctx = GetDeviceOption('cpu') if src_ctx is None: src_ctx = _proto_utils.GetDeviceOption('cpu')
if ctx is None: ctx = GetDeviceOption('cpu') if ctx is None: ctx = _proto_utils.GetDeviceOption('cpu')
dragon.C.TensorFromTensor( _get_workspace().TensorFromTensor(
_stringify_tensor(tensor), _stringify_tensor(src), _stringify_tensor(tensor),
_stringify_proto(ctx), _stringify_proto(src_ctx)) _stringify_tensor(src),
_stringify_proto(ctx),
_stringify_proto(src_ctx))
return tensor return tensor
def FromPyArray(array, name=None): def FromArray(array, name=None):
"""Create a Tensor from a existing Array. """Create a Tensor from a existing Array.
Note that memory of Tensor are ``zero-copied``. Note that memory of Tensor are ``zero-copied``.
...@@ -128,12 +130,13 @@ def FromPyArray(array, name=None): ...@@ -128,12 +130,13 @@ def FromPyArray(array, name=None):
""" """
tensor = _try_get_tensor(name) tensor = _try_get_tensor(name)
if not isinstance(array, numpy.ndarray): if not isinstance(array, numpy.ndarray):
raise TypeError('The given nd-array should be numpy.ndarray.') raise TypeError('Excepted a numpy.ndarray.')
dragon.C.TensorFromPyArray(_stringify_tensor(tensor), array) _get_workspace().TensorFromArray(
_stringify_tensor(tensor), array)
return tensor return tensor
def SetPyArray(tensor, array): def SetArray(tensor, array):
"""Set a Tensor from a existing Array. """Set a Tensor from a existing Array.
Note that memory of Tensor are ``zero-copied``. Note that memory of Tensor are ``zero-copied``.
...@@ -149,15 +152,12 @@ def SetPyArray(tensor, array): ...@@ -149,15 +152,12 @@ def SetPyArray(tensor, array):
------- -------
None None
References
----------
The wrapper of ``TensorFromPyArrayCC``.
""" """
dragon.C.TensorFromPyArray(_stringify_tensor(tensor), array) _get_workspace().TensorFromArray(
_stringify_tensor(tensor), array)
def ToPyArray(tensor, readonly=False): def ToArray(tensor, readonly=False):
"""Create a Array from a existing Tensor. """Create a Array from a existing Tensor.
Note that memory of Array are *zero-copied*. Note that memory of Array are *zero-copied*.
...@@ -175,7 +175,8 @@ def ToPyArray(tensor, readonly=False): ...@@ -175,7 +175,8 @@ def ToPyArray(tensor, readonly=False):
The array sharing the memory with original tensor. The array sharing the memory with original tensor.
""" """
return dragon.C.TensorToPyArray(_stringify_tensor(tensor), readonly) return _get_workspace().TensorToArray(
_stringify_tensor(tensor), readonly)
def GetStorage(tensor): def GetStorage(tensor):
...@@ -193,8 +194,8 @@ def GetStorage(tensor): ...@@ -193,8 +194,8 @@ def GetStorage(tensor):
""" """
tensor = _stringify_tensor(tensor) tensor = _stringify_tensor(tensor)
if not dragon.workspace.HasTensor(tensor): return None if not _get_workspace().HasTensor(tensor): return None
return dragon.C.GetTensor(tensor) return _get_workspace().GetTensor(tensor)
def _stringify_proto(obj): def _stringify_proto(obj):
...@@ -210,5 +211,10 @@ def _stringify_tensor(obj): ...@@ -210,5 +211,10 @@ def _stringify_tensor(obj):
def _try_get_tensor(name=None): def _try_get_tensor(name=None):
"""Try to create or get a tensor""" """Try to create or get a tensor"""
if name is None or name == '': return Tensor() if name is None or name == '': return _Tensor()
else: return Tensor.Ref(name) else: return _Tensor.Ref(name)
\ No newline at end of file
def _get_workspace():
"""Get the current default workspace."""
return _workspace.get_default_workspace()
\ No newline at end of file
Dragon/python/dragon/core/tls.py 0 → 100644
# ------------------------------------------------------------
# Copyright (c) 2017-present, SeetaTech, Co.,Ltd.
#
# Licensed under the BSD 2-Clause License.
# You should have received a copy of the BSD 2-Clause License
# along with the software. If not, See,
#
# <https://opensource.org/licenses/BSD-2-Clause>
#
# ------------------------------------------------------------
"""Define the common thread local structures."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import threading
import contextlib
class Constant(threading.local):
def __init__(self, **attrs):
super(Constant, self).__init__()
self.__dict__.update(attrs)
class Stack(threading.local):
def __init__(self):
super(Stack, self).__init__()
self._enforce_nesting = True
self.stack = []
def get_default(self):
return self.stack[-1] if len(self.stack) >= 1 else None
def reset(self):
self.stack = []
def is_cleared(self):
return not self.stack
@property
def enforce_nesting(self):
return self._enforce_nesting
@enforce_nesting.setter
def enforce_nesting(self, value):
self._enforce_nesting = value
@contextlib.contextmanager
def get_controller(self, default):
"""A context manager for manipulating a default stack."""
self.stack.append(default)
try:
yield default
finally:
# stack may be empty if reset() was called
if self.stack:
if self._enforce_nesting:
if self.stack[-1] is not default:
raise AssertionError(
"Nesting violated for default stack of %s objects" %
type(default))
self.stack.pop()
else:
self.stack.remove(default)
\ No newline at end of file
Dragon/python/dragon/import_c_api.py
...@@ -23,7 +23,6 @@ from __future__ import print_function ...@@ -23,7 +23,6 @@ from __future__ import print_function
import sys import sys
import logging as _logging import logging as _logging
import atexit
try: try:
from dragon.libdragon import * from dragon.libdragon import *
...@@ -32,9 +31,5 @@ except ImportError as e: ...@@ -32,9 +31,5 @@ except ImportError as e:
'Cannot import dragon. Error: {0}'.format(str(e))) 'Cannot import dragon. Error: {0}'.format(str(e)))
sys.exit(1) sys.exit(1)
REGISTERED_OPERATORS = set(s for s in RegisteredOperators()) REGISTERED_OPERATORS = set(s for s in RegisteredOperators())
NO_GRADIENT_OPERATORS = set(s for s in NoGradientOperators()) NO_GRADIENT_OPERATORS = set(s for s in NoGradientOperators())
\ No newline at end of file
atexit.register(OnModuleExit)
\ No newline at end of file
Dragon/python/dragon/memonger.py
...@@ -15,6 +15,8 @@ from __future__ import absolute_import ...@@ -15,6 +15,8 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
from dragon import config as _cfg
def ShareGrads(enabled=True): def ShareGrads(enabled=True):
"""Enable gradients sharing globally. """Enable gradients sharing globally.
...@@ -34,8 +36,8 @@ def ShareGrads(enabled=True): ...@@ -34,8 +36,8 @@ def ShareGrads(enabled=True):
>>> opt.ShareGrads() >>> opt.ShareGrads()
""" """
from dragon.config import option options = _cfg.GetGlobalOptions()
option['share_grads'] = enabled options['share_grads'] = enabled
def IsGradsShared(): def IsGradsShared():
...@@ -47,8 +49,8 @@ def IsGradsShared(): ...@@ -47,8 +49,8 @@ def IsGradsShared():
``True`` if sharing grads else ``False``. ``True`` if sharing grads else ``False``.
""" """
from dragon.config import option options = _cfg.GetGlobalOptions()
return option['share_grads'] return options['share_grads']
def Drop(op_func, *args, **kwargs): def Drop(op_func, *args, **kwargs):
......
Dragon/python/dragon/operators/custom/minibatch.py
...@@ -13,8 +13,8 @@ from __future__ import absolute_import ...@@ -13,8 +13,8 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import dragon from dragon.utils import vision as _vision
import dragon.utils.vision from dragon.core import workspace as _workspace
class MiniBatchOp(object): class MiniBatchOp(object):
...@@ -36,7 +36,7 @@ class MiniBatchOp(object): ...@@ -36,7 +36,7 @@ class MiniBatchOp(object):
""" """
kwargs = eval(self.param_str) kwargs = eval(self.param_str)
self._data_batch = dragon.utils.vision.DataBatch(**kwargs) self._data_batch = _vision.DataBatch(**kwargs)
def run(self, inputs, outputs): def run(self, inputs, outputs):
"""Run method, i.e., forward pass. """Run method, i.e., forward pass.
...@@ -55,4 +55,4 @@ class MiniBatchOp(object): ...@@ -55,4 +55,4 @@ class MiniBatchOp(object):
""" """
blobs = self._data_batch.get() blobs = self._data_batch.get()
for idx, blob in enumerate(blobs): for idx, blob in enumerate(blobs):
dragon.workspace.FeedTensor(outputs[idx], blob) _workspace.FeedTensor(outputs[idx], blob)
\ No newline at end of file \ No newline at end of file
Dragon/python/dragon/ops.py
...@@ -15,149 +15,149 @@ from __future__ import absolute_import ...@@ -15,149 +15,149 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
from .operators import initializer as init_ops from .operators import initializer as _init_ops
from .operators import vision as vision_ops from .operators import vision as _vision_ops
from .operators import loss as loss_ops from .operators import loss as _loss_ops
from .operators import data as data_ops from .operators import data as _data_ops
from .operators import activation as active_ops from .operators import activation as _active_ops
from .operators import arithmetic as math_ops from .operators import arithmetic as _math_ops
from .operators import control_flow as control_flow_ops from .operators import control_flow as _control_flow_ops
from .operators import misc as misc_ops from .operators import misc as _misc_ops
from .operators import mpi as mpi_ops from .operators import mpi as _mpi_ops
from .operators import array as array_ops from .operators import array as _array_ops
from .operators import norm as norm_ops from .operators import norm as _norm_ops
from .operators import recurrent as recurrent_ops from .operators import recurrent as _recurrent_ops
from .operators import contrib as contrib_ops from .operators import contrib as _contrib_ops
# Data # Data
LMDBData = data_ops.LMDBData LMDBData = _data_ops.LMDBData
ImageData = data_ops.ImageData ImageData = _data_ops.ImageData
# Initializer # Initializer
Fill = init_ops.Fill Fill = _init_ops.Fill
RandomUniform = init_ops.RandomUniform RandomUniform = _init_ops.RandomUniform
RandomNormal = init_ops.RandomNormal RandomNormal = _init_ops.RandomNormal
TruncatedNormal = init_ops.TruncatedNormal TruncatedNormal = _init_ops.TruncatedNormal
GlorotUniform = init_ops.GlorotUniform GlorotUniform = _init_ops.GlorotUniform
GlorotNormal = init_ops.GlorotNormal GlorotNormal = _init_ops.GlorotNormal
# Vision # Vision
Conv2d = vision_ops.Conv2d Conv2d = _vision_ops.Conv2d
DepthwiseConv2d = vision_ops.DepthwiseConv2d DepthwiseConv2d = _vision_ops.DepthwiseConv2d
ConvTranspose2d = DeConv2d = Conv2dTranspose = vision_ops.ConvTranspose2d ConvTranspose2d = DeConv2d = Conv2dTranspose = _vision_ops.ConvTranspose2d
Pool2d = vision_ops.Pool2d Pool2d = _vision_ops.Pool2d
ROIPool = vision_ops.ROIPool ROIPool = _vision_ops.ROIPool
ROIAlign = vision_ops.ROIAlign ROIAlign = _vision_ops.ROIAlign
LRN = vision_ops.LRN LRN = _vision_ops.LRN
NNResize = vision_ops.NNResize NNResize = _vision_ops.NNResize
BilinearResize = vision_ops.BilinearResize BilinearResize = _vision_ops.BilinearResize
BiasAdd = vision_ops.BiasAdd BiasAdd = _vision_ops.BiasAdd
DropBlock2d = vision_ops.DropBlock2d DropBlock2d = _vision_ops.DropBlock2d
# Recurrent # Recurrent
LSTMCell = recurrent_ops.LSTMCell LSTMCell = _recurrent_ops.LSTMCell
RNN = recurrent_ops.RNN RNN = _recurrent_ops.RNN
LSTM = recurrent_ops.LSTM LSTM = _recurrent_ops.LSTM
GRU = recurrent_ops.GRU GRU = _recurrent_ops.GRU
# Activation # Activation
Sigmoid = active_ops.Sigmoid Sigmoid = _active_ops.Sigmoid
Tanh = active_ops.Tanh Tanh = _active_ops.Tanh
Relu = active_ops.Relu Relu = _active_ops.Relu
LRelu = active_ops.LRelu LRelu = _active_ops.LRelu
PRelu = active_ops.PRelu PRelu = _active_ops.PRelu
Elu = active_ops.Elu Elu = _active_ops.Elu
SElu = active_ops.SElu SElu = _active_ops.SElu
Softmax = active_ops.Softmax Softmax = _active_ops.Softmax
Dropout = active_ops.Dropout Dropout = _active_ops.Dropout
# Loss # Loss
NLLLoss = loss_ops.NLLLoss NLLLoss = _loss_ops.NLLLoss
SparseSoftmaxCrossEntropy = loss_ops.SparseSoftmaxCrossEntropy SparseSoftmaxCrossEntropy = _loss_ops.SparseSoftmaxCrossEntropy
SigmoidCrossEntropy = loss_ops.SigmoidCrossEntropy SigmoidCrossEntropy = _loss_ops.SigmoidCrossEntropy
SoftmaxCrossEntropy = loss_ops.SoftmaxCrossEntropy SoftmaxCrossEntropy = _loss_ops.SoftmaxCrossEntropy
SmoothL1Loss = loss_ops.SmoothL1Loss SmoothL1Loss = _loss_ops.SmoothL1Loss
L1Loss = loss_ops.L1Loss L1Loss = _loss_ops.L1Loss
L2Loss = loss_ops.L2Loss L2Loss = _loss_ops.L2Loss
SigmoidFocalLoss = loss_ops.SigmoidFocalLoss SigmoidFocalLoss = _loss_ops.SigmoidFocalLoss
SoftmaxFocalLoss = loss_ops.SoftmaxFocalLoss SoftmaxFocalLoss = _loss_ops.SoftmaxFocalLoss
CTCLoss = loss_ops.CTCLoss CTCLoss = _loss_ops.CTCLoss
# Arithmetic # Arithmetic
Add = math_ops.Add Add = _math_ops.Add
Sub = math_ops.Sub Sub = _math_ops.Sub
Mul = math_ops.Mul Mul = _math_ops.Mul
Div = math_ops.Div Div = _math_ops.Div
Maximum = math_ops.Maximum Maximum = _math_ops.Maximum
Minimum = math_ops.Minimum Minimum = _math_ops.Minimum
Moments = math_ops.Moments Moments = _math_ops.Moments
Clip = math_ops.Clip Clip = _math_ops.Clip
Matmul = math_ops.Matmul Matmul = _math_ops.Matmul
Pow = math_ops.Pow Pow = _math_ops.Pow
Dot = math_ops.Dot Dot = _math_ops.Dot
Log = math_ops.Log Log = _math_ops.Log
Exp = math_ops.Exp Exp = _math_ops.Exp
Square = math_ops.Square Square = _math_ops.Square
Sqrt = math_ops.Sqrt Sqrt = _math_ops.Sqrt
FullyConnected = math_ops.FullyConnected FullyConnected = _math_ops.FullyConnected
Eltwise = math_ops.Eltwise Eltwise = _math_ops.Eltwise
Affine = math_ops.Affine Affine = _math_ops.Affine
GramMatrix = math_ops.GramMatrix GramMatrix = _math_ops.GramMatrix
Accumulate = math_ops.Accumulate Accumulate = _math_ops.Accumulate
MovingAverage = math_ops.MovingAverage MovingAverage = _math_ops.MovingAverage
# Normalization # Normalization
BatchNorm = norm_ops.BatchNorm BatchNorm = _norm_ops.BatchNorm
GroupNorm = norm_ops.GroupNorm GroupNorm = _norm_ops.GroupNorm
LayerNorm = norm_ops.LayerNorm LayerNorm = _norm_ops.LayerNorm
InstanceNorm = norm_ops.InstanceNorm InstanceNorm = _norm_ops.InstanceNorm
L2Norm = norm_ops.L2Norm L2Norm = _norm_ops.L2Norm
# NDArray # NDArray
Gather = array_ops.Gather Gather = _array_ops.Gather
Crop = array_ops.Crop Crop = _array_ops.Crop
Reduce = array_ops.Reduce Reduce = _array_ops.Reduce
Sum = array_ops.Sum Sum = _array_ops.Sum
Mean = array_ops.Mean Mean = _array_ops.Mean
Max = array_ops.Max Max = _array_ops.Max
ArgMax = array_ops.ArgMax ArgMax = _array_ops.ArgMax
Min = array_ops.Min Min = _array_ops.Min
ArgMin = array_ops.ArgMin ArgMin = _array_ops.ArgMin
Slice = array_ops.Slice Slice = _array_ops.Slice
Stack = array_ops.Stack Stack = _array_ops.Stack
Concat = array_ops.Concat Concat = _array_ops.Concat
Transpose = array_ops.Transpose Transpose = _array_ops.Transpose
Repeat = array_ops.Repeat Repeat = _array_ops.Repeat
Tile = array_ops.Tile Tile = _array_ops.Tile
Pad = array_ops.Pad Pad = _array_ops.Pad
OneHot = array_ops.OneHot OneHot = _array_ops.OneHot
Flatten = array_ops.Flatten Flatten = _array_ops.Flatten
Reshape = array_ops.Reshape Reshape = _array_ops.Reshape
ExpandDims = array_ops.ExpandDims ExpandDims = _array_ops.ExpandDims
Squeeze = array_ops.Squeeze Squeeze = _array_ops.Squeeze
Shape = array_ops.Shape Shape = _array_ops.Shape
Arange = array_ops.Arange Arange = _array_ops.Arange
Multinomial = array_ops.Multinomial Multinomial = _array_ops.Multinomial
# Control Flow # Control Flow
Copy = control_flow_ops.Copy Copy = _control_flow_ops.Copy
Assign = control_flow_ops.Assign Assign = _control_flow_ops.Assign
Equal = control_flow_ops.Equal Equal = _control_flow_ops.Equal
Less = control_flow_ops.Less Less = _control_flow_ops.Less
LessEqual = control_flow_ops.LessEqual LessEqual = _control_flow_ops.LessEqual
Greater = control_flow_ops.Greater Greater = _control_flow_ops.Greater
GreaterEqual = control_flow_ops.GreaterEqual GreaterEqual = _control_flow_ops.GreaterEqual
# Misc # Misc
Cast = AsType = misc_ops.Cast Cast = AsType = _misc_ops.Cast
Run = misc_ops.Run Run = _misc_ops.Run
Template = misc_ops.Template Template = _misc_ops.Template
Accuracy = misc_ops.Accuracy Accuracy = _misc_ops.Accuracy
StopGradient = misc_ops.StopGradient StopGradient = _misc_ops.StopGradient
# MPI # MPI
MPIBroadcast = mpi_ops.MPIBroadcast MPIBroadcast = _mpi_ops.MPIBroadcast
MPIGather = mpi_ops.MPIGather MPIGather = _mpi_ops.MPIGather
# Contrib # Contrib
Proposal = contrib_ops.Proposal # R-CNN Proposal = _contrib_ops.Proposal # R-CNN
\ No newline at end of file \ No newline at end of file
Dragon/python/dragon/proto/dragon.proto
...@@ -145,18 +145,6 @@ message GradientProto { ...@@ -145,18 +145,6 @@ message GradientProto {
optional string external = 3; optional string external = 3;
} }
// Record the updater information
message UpdaterProto {
// The operator name to use.
optional string name = 1;
// The operator type.
optional string type = 2;
// The tensor to update.
repeated string tensor = 3;
// The arguments.
repeated Argument arg = 4;
}
// Graph Definition // Graph Definition
message GraphDef { message GraphDef {
// The graph name. // The graph name.
...@@ -181,6 +169,4 @@ message GraphDef { ...@@ -181,6 +169,4 @@ message GraphDef {
// The gradients information. // The gradients information.
repeated GradientProto gradient = 9; repeated GradientProto gradient = 9;
// The updaters information.
repeated UpdaterProto updater = 10;
} }
\ No newline at end of file
Dragon/python/dragon/updaters.py
...@@ -22,8 +22,8 @@ from __future__ import print_function ...@@ -22,8 +22,8 @@ from __future__ import print_function
import pprint import pprint
from dragon.core import workspace from dragon.core import workspace as _workspace
from dragon.core.tensor import Tensor from dragon.core.tensor import Tensor as _Tensor
class BaseUpdater(object): class BaseUpdater(object):
...@@ -32,12 +32,14 @@ class BaseUpdater(object): ...@@ -32,12 +32,14 @@ class BaseUpdater(object):
# Store the global unique slot index # Store the global unique slot index
_DEFAULT_UNIQUE_SLOT_ID = 0 _DEFAULT_UNIQUE_SLOT_ID = 0
def __init__(self, def __init__(
scale_gradient=1.0, self,
clip_gradient=-1.0, scale_gradient=1.0,
l2_decay=-1.0, clip_gradient=-1.0,
slot=None, l2_decay=-1.0,
verbose=True): slot=None,
verbose=True,
):
"""Construct a Updater to optimize the objectives. """Construct a Updater to optimize the objectives.
Parameters Parameters
...@@ -84,7 +86,7 @@ class BaseUpdater(object): ...@@ -84,7 +86,7 @@ class BaseUpdater(object):
None None
""" """
pair = (tensor.name if isinstance(tensor, Tensor) \ pair = (tensor.name if isinstance(tensor, _Tensor) \
else tensor for tensor in pair) else tensor for tensor in pair)
self._param_group.append((pair, self._param_group.append((pair,
{'lr_mult': lr_mult, 'decay_mult': decay_mult})) {'lr_mult': lr_mult, 'decay_mult': decay_mult}))
...@@ -93,7 +95,8 @@ class BaseUpdater(object): ...@@ -93,7 +95,8 @@ class BaseUpdater(object):
defaults = self.__dict__.get('_defaults') defaults = self.__dict__.get('_defaults')
if item in defaults: if item in defaults:
if self._registered: if self._registered:
return workspace.FetchTensor(self._slot + '/' + item) return _workspace.FetchTensor(
self._slot + '/' + item)
else: return defaults[item] else: return defaults[item]
return self.__dict__[item] return self.__dict__[item]
...@@ -101,8 +104,9 @@ class BaseUpdater(object): ...@@ -101,8 +104,9 @@ class BaseUpdater(object):
defaults = self.__dict__.get('_defaults') defaults = self.__dict__.get('_defaults')
if defaults is not None and key in defaults: if defaults is not None and key in defaults:
if self._registered: if self._registered:
workspace.FeedTensor(self._slot + '/' + key, value, _workspace.FeedTensor(
dtype='float32', force_cpu=True) self._slot + '/' + key, value,
dtype='float32', force_cpu=True)
else: else:
self._defaults[key] = value self._defaults[key] = value
else: else:
...@@ -111,8 +115,9 @@ class BaseUpdater(object): ...@@ -111,8 +115,9 @@ class BaseUpdater(object):
def register_in_workspace(self): def register_in_workspace(self):
if not self._registered: if not self._registered:
for k, v in self._defaults.items(): for k, v in self._defaults.items():
workspace.FeedTensor(self._slot + "/" + k, v, _workspace.FeedTensor(
dtype='float32', force_cpu=True) self._slot + "/" + k, v,
dtype='float32', force_cpu=True)
self._registered = True self._registered = True
if self._verbose: if self._verbose:
print('---------------------------------------------------------') print('---------------------------------------------------------')
...@@ -206,8 +211,14 @@ class AdamUpdater(BaseUpdater): ...@@ -206,8 +211,14 @@ class AdamUpdater(BaseUpdater):
Introduced by `[Kingma & Ba, 2014] <https://arxiv.org/abs/1412.6980>`_. Introduced by `[Kingma & Ba, 2014] <https://arxiv.org/abs/1412.6980>`_.
""" """
def __init__(self, base_lr=0.01, beta1=0.9, def __init__(
beta2=0.999, eps=1e-8, **kwargs): self,
base_lr=0.01,
beta1=0.9,
beta2=0.999,
eps=1e-8,
**kwargs
):
"""Construct a Adam Updater to optimize the objectives. """Construct a Adam Updater to optimize the objectives.
Parameters Parameters
...@@ -222,7 +233,7 @@ class AdamUpdater(BaseUpdater): ...@@ -222,7 +233,7 @@ class AdamUpdater(BaseUpdater):
The eps. The eps.
""" """
super(AdamUpdater, self).__init__(**kwargs ) super(AdamUpdater, self).__init__(**kwargs)
self._defaults = dict({ self._defaults = dict({
'base_lr': base_lr, 'base_lr': base_lr,
'beta1': beta1, 'beta1': beta1,
......
Dragon/python/dragon/utils/vision/blob_fetcher.py
...@@ -13,11 +13,11 @@ from __future__ import absolute_import ...@@ -13,11 +13,11 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import numpy as np import numpy
from multiprocessing import Process import multiprocessing
class BlobFetcher(Process): class BlobFetcher(multiprocessing.Process):
"""BlobFetcher is deployed to queue blobs from `DataTransformer`_. """BlobFetcher is deployed to queue blobs from `DataTransformer`_.
It is supported to form *NHWC* image blobs and *1d* label blobs. It is supported to form *NHWC* image blobs and *1d* label blobs.
...@@ -37,10 +37,9 @@ class BlobFetcher(Process): ...@@ -37,10 +37,9 @@ class BlobFetcher(Process):
""" """
super(BlobFetcher, self).__init__() super(BlobFetcher, self).__init__()
self._batch_size = kwargs.get('batch_size', 100) self._batch_size = kwargs.get('batch_size', 128)
self._partition = kwargs.get('partition', False) self._partition = kwargs.get('partition', False)
if self._partition: if self._partition: self._batch_size /= kwargs['group_size']
self._batch_size = self._batch_size // kwargs['group_size']
self.Q_in = self.Q_out = None self.Q_in = self.Q_out = None
self.daemon = True self.daemon = True
...@@ -54,9 +53,9 @@ class BlobFetcher(Process): ...@@ -54,9 +53,9 @@ class BlobFetcher(Process):
""" """
im, labels = self.Q_in.get() im, labels = self.Q_in.get()
im_blob = np.zeros(shape=([self._batch_size] + list(im.shape)), dtype=np.uint8) im_blob = numpy.zeros(shape=([self._batch_size] + list(im.shape)), dtype='uint8')
label_blob = np.zeros((self._batch_size, len(labels)), dtype=np.int64) label_blob = numpy.zeros((self._batch_size, len(labels)), dtype='int64')
for ix in range(0, self._batch_size): for ix in range(self._batch_size):
im_blob[ix, :, :, :], label_blob[ix, :] = im, labels im_blob[ix, :, :, :], label_blob[ix, :] = im, labels
if ix != self._batch_size - 1: im, labels = self.Q_in.get() if ix != self._batch_size - 1: im, labels = self.Q_in.get()
return im_blob, label_blob return im_blob, label_blob
......
Dragon/python/dragon/utils/vision/data_batch.py
...@@ -14,11 +14,10 @@ from __future__ import division ...@@ -14,11 +14,10 @@ from __future__ import division
from __future__ import print_function from __future__ import print_function
import time import time
import pprint import multiprocessing
from multiprocessing import Queue
import dragon.core.mpi as mpi from dragon.core import mpi as _mpi
import dragon.core.logging as logging from dragon.core import logging as _logging
from .data_reader import DataReader from .data_reader import DataReader
from .data_transformer import DataTransformer from .data_transformer import DataTransformer
...@@ -77,10 +76,11 @@ class DataBatch(object): ...@@ -77,10 +76,11 @@ class DataBatch(object):
super(DataBatch, self).__init__() super(DataBatch, self).__init__()
# Init mpi # Init mpi
global_rank = 0; local_rank = 0; group_size = 1 global_rank = 0; local_rank = 0; group_size = 1
if mpi.Is_Init(): if _mpi.Is_Init() and kwargs.get(
idx, group = mpi.AllowParallel() 'phase', 'TRAIN') == 'TRAIN':
if idx != -1: # DataParallel rank, group = _mpi.AllowParallel()
global_rank = mpi.Rank() if rank != -1: # DataParallel
global_rank = _mpi.Rank()
group_size = len(group) group_size = len(group)
for i, node in enumerate(group): for i, node in enumerate(group):
if global_rank == node: local_rank = i if global_rank == node: local_rank = i
...@@ -105,7 +105,7 @@ class DataBatch(object): ...@@ -105,7 +105,7 @@ class DataBatch(object):
self._num_transformers += 1 self._num_transformers += 1
# Add 1 transformer for random crop # Add 1 transformer for random crop
if kwargs.get('crop_size', 0) > 0 and \ if kwargs.get('crop_size', 0) > 0 and \
kwargs.get('phase', 'TEST') == 'TRAIN': kwargs.get('phase', 'TRAIN') == 'TRAIN':
self._num_transformers += 1 self._num_transformers += 1
self._num_transformers = min(self._num_transformers, self._max_transformers) self._num_transformers = min(self._num_transformers, self._max_transformers)
...@@ -115,9 +115,12 @@ class DataBatch(object): ...@@ -115,9 +115,12 @@ class DataBatch(object):
self._batch_size = int(self._batch_size / kwargs['group_size']) self._batch_size = int(self._batch_size / kwargs['group_size'])
# Init queues # Init queues
self.Q_level_1 = Queue(self._prefetch * self._num_readers * self._batch_size) self.Q_level_1 = multiprocessing.Queue(
self.Q_level_2 = Queue(self._prefetch * self._num_readers * self._batch_size) self._prefetch * self._num_readers * self._batch_size)
self.Q_level_3 = Queue(self._prefetch * self._num_readers) self.Q_level_2 = multiprocessing.Queue(
self._prefetch * self._num_readers * self._batch_size)
self.Q_level_3 = multiprocessing.Queue(
self._prefetch * self._num_readers)
# Init readers # Init readers
self._readers = [] self._readers = []
...@@ -167,11 +170,11 @@ class DataBatch(object): ...@@ -167,11 +170,11 @@ class DataBatch(object):
process.terminate() process.terminate()
process.join() process.join()
terminate(self._fetchers) terminate(self._fetchers)
if local_rank == 0: logging.info('Terminating BlobFetcher ......') if local_rank == 0: _logging.info('Terminate BlobFetcher.')
terminate(self._transformers) terminate(self._transformers)
if local_rank == 0: logging.info('Terminating DataTransformer ......') if local_rank == 0: _logging.info('Terminate DataTransformer.')
terminate(self._readers) terminate(self._readers)
if local_rank == 0: logging.info('Terminating DataReader......') if local_rank == 0: _logging.info('Terminate DataReader.')
import atexit import atexit
atexit.register(cleanup) atexit.register(cleanup)
......
Dragon/python/dragon/utils/vision/data_reader.py
...@@ -14,15 +14,14 @@ from __future__ import division ...@@ -14,15 +14,14 @@ from __future__ import division
from __future__ import print_function from __future__ import print_function
import math import math
import numpy as np import numpy
import numpy.random as npr import multiprocessing
from multiprocessing import Process
import dragon.config as config from dragon import config as _cfg
from dragon.tools.db import LMDB from dragon.tools import db as _db
class DataReader(Process): class DataReader(multiprocessing.Process):
"""DataReader is deployed to queue encoded str from `LMDB`_. """DataReader is deployed to queue encoded str from `LMDB`_.
It is supported to adaptively partition and shuffle records over all distributed nodes. It is supported to adaptively partition and shuffle records over all distributed nodes.
...@@ -55,7 +54,7 @@ class DataReader(Process): ...@@ -55,7 +54,7 @@ class DataReader(Process):
self._part_idx, self._num_parts = 0, 1 self._part_idx, self._num_parts = 0, 1
self._cur_idx, self._cur_chunk_idx = 0, 0 self._cur_idx, self._cur_chunk_idx = 0, 0
self._random_seed = config.GetRandomSeed() self._random_seed = _cfg.GetRandomSeed()
self.Q_out = None self.Q_out = None
self.daemon = True self.daemon = True
...@@ -106,7 +105,9 @@ class DataReader(Process): ...@@ -106,7 +105,9 @@ class DataReader(Process):
""" """
if self._multiple_nodes or self._use_shuffle: if self._multiple_nodes or self._use_shuffle:
if self._use_shuffle: self._perm = npr.permutation(self._num_shuffle_parts) if self._use_shuffle:
self._perm = numpy.random.permutation(
self._num_shuffle_parts)
self._cur_chunk_idx = 0 self._cur_chunk_idx = 0
self._start_idx = int(self._part_idx * self._num_shuffle_parts + self._perm[self._cur_chunk_idx]) self._start_idx = int(self._part_idx * self._num_shuffle_parts + self._perm[self._cur_chunk_idx])
self._start_idx = int(self._start_idx * self._chunk_size) self._start_idx = int(self._start_idx * self._chunk_size)
...@@ -158,23 +159,23 @@ class DataReader(Process): ...@@ -158,23 +159,23 @@ class DataReader(Process):
""" """
# fix seed # fix seed
npr.seed(self._random_seed) numpy.random.seed(self._random_seed)
# init db # init db
self._db = LMDB() self._db = _db.LMDB()
self._db.open(self._source) self._db.open(self._source)
self._zfill = self._db.zfill() self._zfill = self._db.zfill()
self._num_entries = self._db.num_entries() self._num_entries = self._db.num_entries()
self._epoch_size = int(self._num_entries/ self._num_parts + 1) self._epoch_size = int(self._num_entries / self._num_parts + 1)
if self._use_shuffle: if self._use_shuffle:
if self._chunk_size == 1: if self._chunk_size == 1:
# Each chunk has at most 1 record [For Fully Shuffle] # Each chunk has at most 1 record (Naive Shuffle)
self._chunk_size, self._num_shuffle_parts = \ self._chunk_size, self._num_shuffle_parts = \
1, int(self._num_entries / self._num_parts) + 1 1, int(self._num_entries / self._num_parts) + 1
else: else:
if self._use_shuffle and self._chunk_size == -1: if self._use_shuffle and self._chunk_size == -1:
# Search a optimal chunk size by chunks [For Chunk Shuffle] # Search a optimal chunk size by chunks (Chunk Shuffle)
max_chunk_size = self._db._total_size / ((self._num_chunks * (1 << 20))) max_chunk_size = self._db._total_size / ((self._num_chunks * (1 << 20)))
min_chunk_size = 1 min_chunk_size = 1
while min_chunk_size * 2 < max_chunk_size: min_chunk_size *= 2 while min_chunk_size * 2 < max_chunk_size: min_chunk_size *= 2
...@@ -184,17 +185,17 @@ class DataReader(Process): ...@@ -184,17 +185,17 @@ class DataReader(Process):
self._chunk_size = int(self._num_entries / self._num_shuffle_parts / self._num_parts + 1) self._chunk_size = int(self._num_entries / self._num_shuffle_parts / self._num_parts + 1)
limit = (self._num_parts - 0.5) * self._num_shuffle_parts * self._chunk_size limit = (self._num_parts - 0.5) * self._num_shuffle_parts * self._chunk_size
if self._num_entries <= limit: if self._num_entries <= limit:
# Roll back to fully shuffle # Roll back to naive shuffle
self._chunk_size, self._num_shuffle_parts = \ self._chunk_size, self._num_shuffle_parts = \
1, int(self._num_entries / self._num_parts) + 1 1, int(self._num_entries / self._num_parts) + 1
else: else:
# Each chunk has at most K records [For Multiple Nodes] # Each chunk has at most K records
# Note that if ``shuffle`` and ``multiple_nodes`` are all ``False``, # Note that if ``shuffle`` and ``multiple_nodes`` are all *False*,
# ``chunk_size`` and ``num_shuffle_parts`` are meaningless # ``chunk_size`` and ``num_shuffle_parts`` are meaningless
self._chunk_size = int(self._num_entries / self._num_parts) + 1 self._chunk_size = int(self._num_entries / self._num_parts) + 1
self._num_shuffle_parts = 1 self._num_shuffle_parts = 1
self._perm = np.arange(self._num_shuffle_parts) self._perm = numpy.arange(self._num_shuffle_parts)
# Init env # Init env
self.reset() self.reset()
......
Dragon/python/dragon/utils/vision/data_transformer.py
...@@ -13,12 +13,11 @@ from __future__ import absolute_import ...@@ -13,12 +13,11 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import numpy as np import numpy
import numpy.random as npr import multiprocessing
from multiprocessing import Process
import dragon.config as config from dragon import config as _cfg
import dragon.vm.caffe.proto.caffe_pb2 as pb from dragon.vm.caffe.proto import caffe_pb2 as _proto_def
try: try:
import cv2 import cv2
...@@ -31,7 +30,7 @@ except ImportError as e: ...@@ -31,7 +30,7 @@ except ImportError as e:
print("Failed to import PIL. \nIt's OK if disabling color augmentation.".format(str(e))) print("Failed to import PIL. \nIt's OK if disabling color augmentation.".format(str(e)))
class DataTransformer(Process): class DataTransformer(multiprocessing.Process):
"""DataTransformer is deployed to queue transformed images from `DataReader`_. """DataTransformer is deployed to queue transformed images from `DataReader`_.
Nearly all common image augmentation methods are supported. Nearly all common image augmentation methods are supported.
...@@ -72,7 +71,7 @@ class DataTransformer(Process): ...@@ -72,7 +71,7 @@ class DataTransformer(Process):
self._max_random_scale = kwargs.get('max_random_scale', 1.0) self._max_random_scale = kwargs.get('max_random_scale', 1.0)
self._force_color = kwargs.get('force_color', False) self._force_color = kwargs.get('force_color', False)
self._phase = kwargs.get('phase', 'TRAIN') self._phase = kwargs.get('phase', 'TRAIN')
self._random_seed = config.GetRandomSeed() self._random_seed = _cfg.GetRandomSeed()
self.Q_in = self.Q_out = None self.Q_in = self.Q_out = None
self.daemon = True self.daemon = True
...@@ -91,16 +90,16 @@ class DataTransformer(Process): ...@@ -91,16 +90,16 @@ class DataTransformer(Process):
""" """
# decode # decode
datum = pb.Datum() datum = _proto_def.Datum()
datum.ParseFromString(serialized) datum.ParseFromString(serialized)
im = np.fromstring(datum.data, np.uint8) im = numpy.fromstring(datum.data, numpy.uint8)
if datum.encoded is True: if datum.encoded is True:
im = cv2.imdecode(im, -1) im = cv2.imdecode(im, -1)
else: else:
im = im.reshape((datum.height, datum.width, datum.channels)) im = im.reshape((datum.height, datum.width, datum.channels))
# Random scale # Random scale
random_scale = npr.uniform() * ( random_scale = numpy.random.uniform() * (
self._max_random_scale - self._min_random_scale) \ self._max_random_scale - self._min_random_scale) \
+ self._min_random_scale + self._min_random_scale
if random_scale != 1.0: if random_scale != 1.0:
...@@ -109,7 +108,7 @@ class DataTransformer(Process): ...@@ -109,7 +108,7 @@ class DataTransformer(Process):
# Padding # Padding
if self._padding > 0: if self._padding > 0:
pad_img = np.empty(( pad_img = numpy.empty((
im.shape[0] + 2 * self._padding, im.shape[0] + 2 * self._padding,
im.shape[1] + 2 * self._padding, im.shape[2]), dtype=im.dtype) im.shape[1] + 2 * self._padding, im.shape[2]), dtype=im.dtype)
pad_img.fill(self._fill_value) pad_img.fill(self._fill_value)
...@@ -120,8 +119,8 @@ class DataTransformer(Process): ...@@ -120,8 +119,8 @@ class DataTransformer(Process):
# Random crop # Random crop
if self._crop_size > 0: if self._crop_size > 0:
if self._phase == 'TRAIN': if self._phase == 'TRAIN':
h_off = npr.randint(im.shape[0] - self._crop_size + 1) h_off = numpy.random.randint(im.shape[0] - self._crop_size + 1)
w_off = npr.randint(im.shape[1] - self._crop_size + 1) w_off = numpy.random.randint(im.shape[1] - self._crop_size + 1)
else: else:
h_off = int((im.shape[0] - self._crop_size) / 2) h_off = int((im.shape[0] - self._crop_size) / 2)
w_off = int((im.shape[1] - self._crop_size) / 2) w_off = int((im.shape[1] - self._crop_size) / 2)
...@@ -130,28 +129,28 @@ class DataTransformer(Process): ...@@ -130,28 +129,28 @@ class DataTransformer(Process):
# Random mirror # Random mirror
if self._mirror: if self._mirror:
if npr.randint(0, 2) > 0: if numpy.random.randint(0, 2) > 0:
im = im[:, ::-1, :] im = im[:, ::-1, :]
# Gray Transformation # Gray Transformation
if self._force_color: if self._force_color:
if im.shape[2] == 1: if im.shape[2] == 1:
# duplicate to 3 channels # duplicate to 3 channels
im = np.concatenate([im, im, im], axis=2) im = numpy.concatenate([im, im, im], axis=2)
# Color Augmentation # Color Augmentation
if self._color_aug: if self._color_aug:
im = PIL.Image.fromarray(im) im = PIL.Image.fromarray(im)
delta_brightness = npr.uniform(-0.4, 0.4) + 1.0 delta_brightness = numpy.random.uniform(-0.4, 0.4) + 1.0
delta_contrast = npr.uniform(-0.4, 0.4) + 1.0 delta_contrast = numpy.random.uniform(-0.4, 0.4) + 1.0
delta_saturation = npr.uniform(-0.4, 0.4) + 1.0 delta_saturation = numpy.random.uniform(-0.4, 0.4) + 1.0
im = PIL.ImageEnhance.Brightness(im) im = PIL.ImageEnhance.Brightness(im)
im = im.enhance(delta_brightness) im = im.enhance(delta_brightness)
im = PIL.ImageEnhance.Contrast(im) im = PIL.ImageEnhance.Contrast(im)
im = im.enhance(delta_contrast) im = im.enhance(delta_contrast)
im = PIL.ImageEnhance.Color(im) im = PIL.ImageEnhance.Color(im)
im = im.enhance(delta_saturation) im = im.enhance(delta_saturation)
im = np.array(im) im = numpy.array(im)
# Extract Labels # Extract Labels
labels = [] labels = []
...@@ -169,7 +168,7 @@ class DataTransformer(Process): ...@@ -169,7 +168,7 @@ class DataTransformer(Process):
""" """
# Fix the random seed # Fix the random seed
npr.seed(self._random_seed) numpy.random.seed(self._random_seed)
# Run! # Run!
while True: while True:
......
Dragon/python/dragon/utils/vision/im2db.py
...@@ -16,8 +16,8 @@ import shutil ...@@ -16,8 +16,8 @@ import shutil
import argparse import argparse
import cv2 import cv2
from dragon.tools.db import LMDB from dragon.tools import db as _db
from dragon.vm.caffe.proto import caffe_pb2 from dragon.vm.caffe.proto import caffe_pb2 as _proto_def
def resize_image(im, resize): def resize_image(im, resize):
...@@ -37,11 +37,10 @@ def resize_image(im, resize): ...@@ -37,11 +37,10 @@ def resize_image(im, resize):
""" """
if im.shape[0] > im.shape[1]: if im.shape[0] > im.shape[1]:
newsize = (resize, im.shape[0] * resize / im.shape[1]) new_size = (resize, im.shape[0] * resize // im.shape[1])
else: else:
newsize = (im.shape[1] * resize / im.shape[0], resize) new_size = (im.shape[1] * resize // im.shape[0], resize)
im = cv2.resize(im, newsize) return cv2.resize(im, new_size, interpolation=cv2.INTER_LINEAR)
return im
def make_db(args): def make_db(args):
...@@ -72,7 +71,7 @@ def make_db(args): ...@@ -72,7 +71,7 @@ def make_db(args):
print('start time: ', time.strftime("%a, %d %b %Y %H:%M:%S", time.gmtime())) print('start time: ', time.strftime("%a, %d %b %Y %H:%M:%S", time.gmtime()))
db = LMDB(max_commit=10000) db = _db.LMDB(max_commit=10000)
db.open(args.database, mode='w') db.open(args.database, mode='w')
total_line = sum(1 for line in open(args.list)) total_line = sum(1 for line in open(args.list))
...@@ -106,7 +105,7 @@ def make_db(args): ...@@ -106,7 +105,7 @@ def make_db(args):
img = resize_image(img, args.resize) img = resize_image(img, args.resize)
result, imgencode = cv2.imencode('.jpg', img, encode_param) result, imgencode = cv2.imencode('.jpg', img, encode_param)
datum = caffe_pb2.Datum() datum = _proto_def.Datum()
datum.height, datum.width, datum.channels = img.shape datum.height, datum.width, datum.channels = img.shape
datum.label = int(label) datum.label = int(label)
datum.encoded = True datum.encoded = True
......
Dragon/python/dragon/vm/caffe/layer.py
...@@ -15,7 +15,8 @@ from __future__ import absolute_import ...@@ -15,7 +15,8 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import dragon from dragon.core import scope as _scope
from dragon.core.tensor import Tensor as _Tensor
class Layer(object): class Layer(object):
...@@ -74,12 +75,12 @@ class Layer(object): ...@@ -74,12 +75,12 @@ class Layer(object):
# Note that a non-empty tensor scope will make it # Note that a non-empty tensor scope will make it
# impossible to load/save caffe models. You should use # impossible to load/save caffe models. You should use
# a new workspace instead of the terrible name scope # a new workspace instead of the terrible name scope
scoped_name = dragon.get_default_name_scope() + self._name scoped_name = _scope.get_default_name_scope() + self._name
param_name = scoped_name + '/param:{}'.format(len(self._blobs)) param_name = scoped_name + '/param:{}'.format(len(self._blobs))
# Set the name explicitly # Set the name explicitly
variable = dragon.Tensor.Ref(param_name) variable = _Tensor.Ref(param_name)
variable_grad = dragon.Tensor.Ref(param_name + '_grad') variable_grad = _Tensor.Ref(param_name + '_grad')
if filler is not None: if filler is not None:
variable.Fill(**filler) variable.Fill(**filler)
......
Dragon/python/dragon/vm/caffe/layers/data.py
...@@ -15,13 +15,12 @@ from __future__ import absolute_import ...@@ -15,13 +15,12 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import dragon from dragon import ops as _ops
from ..layer import Layer from ..layer import Layer as _Layer
class DataLayer(Layer): class DataLayer(_Layer):
""" """The implementation of ``DataLayer``.
The implementation of ``DataLayer``.
Different from ``Caffe``, we force to use `LMDB`_ backend. Different from ``Caffe``, we force to use `LMDB`_ backend.
...@@ -33,7 +32,7 @@ class DataLayer(Layer): ...@@ -33,7 +32,7 @@ class DataLayer(Layer):
The prefetch count. Refer `DataParameter.prefetch`_. The prefetch count. Refer `DataParameter.prefetch`_.
batch_size : int batch_size : int
The size of a mini-batch. Refer `DataParameter.batch_size`_. The size of a mini-batch. Refer `DataParameter.batch_size`_.
phase : caffe_pb2.Phase phase : Phase
The phase of layer. Refer `LayerParameter.phase`_. The phase of layer. Refer `LayerParameter.phase`_.
mirrow : boolean mirrow : boolean
Whether to randomly mirror. Refer `TransformationParameter.mirror`_. Whether to randomly mirror. Refer `TransformationParameter.mirror`_.
...@@ -49,9 +48,9 @@ class DataLayer(Layer): ...@@ -49,9 +48,9 @@ class DataLayer(Layer):
The min scale of the images. Extension of `TransformationParameter`_. The min scale of the images. Extension of `TransformationParameter`_.
max_random_scale : float max_random_scale : float
The max scale of the images. Extension of `TransformationParameter`_. The max scale of the images. Extension of `TransformationParameter`_.
dtype : caffe_pb2.MemoryDataParameter.DataType dtype : MemoryDataParameter.DataType
The output data type. ``FLOAT32`` or ``FLOAT16``. The output data type. *FLOAT32* or *FLOAT16*.
mean_value : list of float mean_value : sequence of float
The mean of each channel. Refer `TransformationParameter.mean_value`_. The mean of each channel. Refer `TransformationParameter.mean_value`_.
scale : float scale : float
The scaling factor. Refer `TransformationParameter.scale`_. The scaling factor. Refer `TransformationParameter.scale`_.
...@@ -93,20 +92,20 @@ class DataLayer(Layer): ...@@ -93,20 +92,20 @@ class DataLayer(Layer):
[1. / transform_param.scale] * 3 [1. / transform_param.scale] * 3
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
data, label = dragon.ops.LMDBData(**self.arguments) data, label = _ops.LMDBData(**self.arguments)
return dragon.ops.ImageData(data, **self.arguments), label return _ops.ImageData(data, **self.arguments), label
class MemoryDataLayer(Layer): class MemoryDataLayer(_Layer):
"""The implementation of ``MemoryDataLayer``. """The implementation of ``MemoryDataLayer``.
We extend it with ``FP16`` and ``NHWC => NCHW``. We extend it with ``FP16`` and ``NHWC => NCHW``.
Parameters Parameters
---------- ----------
dtype : caffe_pb2.MemoryDataParameter.DataType dtype : MemoryDataParameter.DataType
The output data type. ``FLOAT32`` or ``FLOAT16``. The output data type. ``FLOAT32`` or ``FLOAT16``.
mean_value : list of float mean_value : sequence of float
The mean of each channel. Refer `TransformationParameter.mean_value`_. The mean of each channel. Refer `TransformationParameter.mean_value`_.
scale : float scale : float
The scaling factor. Refer `TransformationParameter.scale`_. The scaling factor. Refer `TransformationParameter.scale`_.
...@@ -131,4 +130,4 @@ class MemoryDataLayer(Layer): ...@@ -131,4 +130,4 @@ class MemoryDataLayer(Layer):
[1. / transform_param.scale] * 3 [1. / transform_param.scale] * 3
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
return dragon.ops.ImageData(bottom, **self.arguments) return _ops.ImageData(bottom, **self.arguments)
\ No newline at end of file \ No newline at end of file
Dragon/python/dragon/vm/caffe/layers/loss.py
...@@ -15,11 +15,11 @@ from __future__ import absolute_import ...@@ -15,11 +15,11 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import dragon from dragon import ops as _ops
from ..layer import Layer from ..layer import Layer as _Layer
class SoftmaxWithLossLayer(Layer): class SoftmaxWithLossLayer(_Layer):
"""The implementation of ``SoftmaxWithLossLayer``. """The implementation of ``SoftmaxWithLossLayer``.
Parameters Parameters
...@@ -52,12 +52,12 @@ class SoftmaxWithLossLayer(Layer): ...@@ -52,12 +52,12 @@ class SoftmaxWithLossLayer(Layer):
} }
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
loss = dragon.ops.SparseSoftmaxCrossEntropy(bottom, **self.arguments) loss = _ops.SparseSoftmaxCrossEntropy(bottom, **self.arguments)
if self._loss_weight is not None: loss *= self._loss_weight if self._loss_weight is not None: loss *= self._loss_weight
return loss return loss
class SigmoidCrossEntropyLossLayer(Layer): class SigmoidCrossEntropyLossLayer(_Layer):
"""The implementation of ``SigmoidCrossEntropyLossLayer``. """The implementation of ``SigmoidCrossEntropyLossLayer``.
Parameters Parameters
...@@ -79,12 +79,12 @@ class SigmoidCrossEntropyLossLayer(Layer): ...@@ -79,12 +79,12 @@ class SigmoidCrossEntropyLossLayer(Layer):
self.arguments = {'normalization': normalization} self.arguments = {'normalization': normalization}
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
loss = dragon.ops.SigmoidCrossEntropy(bottom, **self.arguments) loss = _ops.SigmoidCrossEntropy(bottom, **self.arguments)
if self._loss_weight is not None: loss *= self._loss_weight if self._loss_weight is not None: loss *= self._loss_weight
return loss return loss
class L2LossLayer(Layer): class L2LossLayer(_Layer):
"""The implementation of ``L2LossLayer``. """The implementation of ``L2LossLayer``.
Parameters Parameters
...@@ -106,12 +106,12 @@ class L2LossLayer(Layer): ...@@ -106,12 +106,12 @@ class L2LossLayer(Layer):
self.arguments = {'normalization': normalization} self.arguments = {'normalization': normalization}
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
loss = dragon.ops.L2Loss(bottom, **self.arguments) loss = _ops.L2Loss(bottom, **self.arguments)
if self._loss_weight is not None: loss *= self._loss_weight if self._loss_weight is not None: loss *= self._loss_weight
return loss return loss
class SmoothL1LossLayer(Layer): class SmoothL1LossLayer(_Layer):
"""The implementation of ``SmoothL1LossLayer``. """The implementation of ``SmoothL1LossLayer``.
Parameters Parameters
...@@ -140,12 +140,12 @@ class SmoothL1LossLayer(Layer): ...@@ -140,12 +140,12 @@ class SmoothL1LossLayer(Layer):
} }
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
loss = dragon.ops.SmoothL1Loss(bottom, **self.arguments) loss = _ops.SmoothL1Loss(bottom, **self.arguments)
if self._loss_weight is not None: loss *= self._loss_weight if self._loss_weight is not None: loss *= self._loss_weight
return loss return loss
class SigmoidWithFocalLossLayer(Layer): class SigmoidWithFocalLossLayer(_Layer):
"""The implementation of ``SigmoidWithFocalLossLayer``. """The implementation of ``SigmoidWithFocalLossLayer``.
Parameters Parameters
...@@ -183,12 +183,12 @@ class SigmoidWithFocalLossLayer(Layer): ...@@ -183,12 +183,12 @@ class SigmoidWithFocalLossLayer(Layer):
} }
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
loss = dragon.ops.SigmoidFocalLoss(bottom, **self.arguments) loss = _ops.SigmoidFocalLoss(bottom, **self.arguments)
if self._loss_weight is not None: loss *= self._loss_weight if self._loss_weight is not None: loss *= self._loss_weight
return loss return loss
class SoftmaxWithFocalLossLayer(Layer): class SoftmaxWithFocalLossLayer(_Layer):
"""The implementation of ``SoftmaxWithFocalLossLayer``. """The implementation of ``SoftmaxWithFocalLossLayer``.
Parameters Parameters
...@@ -227,6 +227,6 @@ class SoftmaxWithFocalLossLayer(Layer): ...@@ -227,6 +227,6 @@ class SoftmaxWithFocalLossLayer(Layer):
} }
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
loss = dragon.ops.SoftmaxFocalLoss(bottom, **self.arguments) loss = _ops.SoftmaxFocalLoss(bottom, **self.arguments)
if self._loss_weight is not None: loss *= self._loss_weight if self._loss_weight is not None: loss *= self._loss_weight
return loss return loss
\ No newline at end of file
Dragon/python/dragon/vm/caffe/layers/mpi.py
...@@ -15,11 +15,11 @@ from __future__ import absolute_import ...@@ -15,11 +15,11 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import dragon from dragon import ops as _ops
from ..layer import Layer from ..layer import Layer as _Layer
class MPIBroadcastLayer(Layer): class MPIBroadcastLayer(_Layer):
"""The implementation of ``MPIBroadcastLayer``. """The implementation of ``MPIBroadcastLayer``.
Parameters Parameters
...@@ -33,10 +33,10 @@ class MPIBroadcastLayer(Layer): ...@@ -33,10 +33,10 @@ class MPIBroadcastLayer(Layer):
self.arguments = {'root': LayerParameter.mpi_param.root} self.arguments = {'root': LayerParameter.mpi_param.root}
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
return dragon.ops.MPIBroadcast(bottom, **self.arguments) return _ops.MPIBroadcast(bottom, **self.arguments)
class MPIGatherLayer(Layer): class MPIGatherLayer(_Layer):
"""The implementation of ``MPIGatherLayer``. """The implementation of ``MPIGatherLayer``.
Parameters Parameters
...@@ -53,4 +53,4 @@ class MPIGatherLayer(Layer): ...@@ -53,4 +53,4 @@ class MPIGatherLayer(Layer):
} }
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
return dragon.ops.MPIGather(bottom, **self.arguments) return _ops.MPIGather(bottom, **self.arguments)
\ No newline at end of file \ No newline at end of file
Dragon/python/dragon/vm/caffe/layers/neuron.py
...@@ -15,11 +15,11 @@ from __future__ import absolute_import ...@@ -15,11 +15,11 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import dragon from dragon import ops as _ops
from ..layer import Layer from ..layer import Layer as _Layer
class ReLULayer(Layer): class ReLULayer(_Layer):
"""The implementation of ``ReLULayer``. """The implementation of ``ReLULayer``.
Parameters Parameters
...@@ -35,10 +35,10 @@ class ReLULayer(Layer): ...@@ -35,10 +35,10 @@ class ReLULayer(Layer):
self.arguments = {'slope': param.negative_slope} self.arguments = {'slope': param.negative_slope}
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
return dragon.ops.Relu(bottom, **self.arguments) return _ops.Relu(bottom, **self.arguments)
class PReLULayer(Layer): class PReLULayer(_Layer):
"""The implementation of ``PReLULayer``. """The implementation of ``PReLULayer``.
Parameters Parameters
...@@ -61,10 +61,10 @@ class PReLULayer(Layer): ...@@ -61,10 +61,10 @@ class PReLULayer(Layer):
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
inputs = [bottom] + [blob['data'] for blob in self._blobs] inputs = [bottom] + [blob['data'] for blob in self._blobs]
return dragon.ops.PRelu(inputs, **self.arguments) return _ops.PRelu(inputs, **self.arguments)
class ELULayer(Layer): class ELULayer(_Layer):
"""The implementation of ``ELULayer``. """The implementation of ``ELULayer``.
Parameters Parameters
...@@ -78,40 +78,40 @@ class ELULayer(Layer): ...@@ -78,40 +78,40 @@ class ELULayer(Layer):
self.arguments = {'alpha': float(LayerParameter.elu_param.alpha)} self.arguments = {'alpha': float(LayerParameter.elu_param.alpha)}
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
return dragon.ops.Elu(bottom, **self.arguments) return _ops.Elu(bottom, **self.arguments)
class SELULayer(Layer): class SELULayer(_Layer):
"""The implementation of ``SELULayer``.""" """The implementation of ``SELULayer``."""
def __init__(self, LayerParameter): def __init__(self, LayerParameter):
super(SELULayer, self).__init__(LayerParameter) super(SELULayer, self).__init__(LayerParameter)
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
return dragon.ops.SElu(bottom, **self.arguments) return _ops.SElu(bottom, **self.arguments)
class SigmoidLayer(Layer): class SigmoidLayer(_Layer):
"""The implementation of ``SigmoidLayer``.""" """The implementation of ``SigmoidLayer``."""
def __init__(self, LayerParameter): def __init__(self, LayerParameter):
super(SigmoidLayer, self).__init__(LayerParameter) super(SigmoidLayer, self).__init__(LayerParameter)
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
return dragon.ops.Sigmoid(bottom, **self.arguments) return _ops.Sigmoid(bottom, **self.arguments)
class TanHLayer(Layer): class TanHLayer(_Layer):
"""The implementation of ``TanHLayer``.""" """The implementation of ``TanHLayer``."""
def __init__(self, LayerParameter): def __init__(self, LayerParameter):
super(TanHLayer, self).__init__(LayerParameter) super(TanHLayer, self).__init__(LayerParameter)
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
return dragon.ops.Tanh(bottom, **self.arguments) return _ops.Tanh(bottom, **self.arguments)
class DropoutLayer(Layer): class DropoutLayer(_Layer):
"""The implementation of ``DropoutLayer``. """The implementation of ``DropoutLayer``.
Parameters Parameters
...@@ -132,10 +132,10 @@ class DropoutLayer(Layer): ...@@ -132,10 +132,10 @@ class DropoutLayer(Layer):
} }
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
return dragon.ops.Dropout(bottom, **self.arguments) return _ops.Dropout(bottom, **self.arguments)
class PowerLayer(Layer): class PowerLayer(_Layer):
"""The implementation of ``PowerLayer``. """The implementation of ``PowerLayer``.
Parameters Parameters
...@@ -158,4 +158,4 @@ class PowerLayer(Layer): ...@@ -158,4 +158,4 @@ class PowerLayer(Layer):
} }
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
return dragon.ops.Pow(bottom, **self.arguments) return _ops.Pow(bottom, **self.arguments)
\ No newline at end of file \ No newline at end of file
Dragon/python/dragon/vm/caffe/layers/vision.py
...@@ -15,11 +15,11 @@ from __future__ import absolute_import ...@@ -15,11 +15,11 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import dragon from dragon import ops as _ops
from ..layer import Layer from ..layer import Layer as _Layer
class ConvolutionLayer(Layer): class ConvolutionLayer(_Layer):
"""The implementation of ``ConvolutionLayer``. """The implementation of ``ConvolutionLayer``.
Parameters Parameters
...@@ -28,19 +28,19 @@ class ConvolutionLayer(Layer): ...@@ -28,19 +28,19 @@ class ConvolutionLayer(Layer):
The output channels. Refer `ConvolutionParameter.num_output`_. The output channels. Refer `ConvolutionParameter.num_output`_.
bias_term : boolean bias_term : boolean
Whether to use bias. Refer `ConvolutionParameter.bias_term`_. Whether to use bias. Refer `ConvolutionParameter.bias_term`_.
pad : list of int pad : sequence of int
The zero padding size(s). Refer `ConvolutionParameter.pad`_. The zero padding size(s). Refer `ConvolutionParameter.pad`_.
kernel_size : list of int kernel_size : list of int
The kernel size(s). Refer `ConvolutionParameter.kernel_size`_. The kernel size(s). Refer `ConvolutionParameter.kernel_size`_.
stride : list of int stride : sequence of int
The stride(s). Refer `ConvolutionParameter.stride`_. The stride(s). Refer `ConvolutionParameter.stride`_.
dilation : list of int dilation : sequence of int
The dilation(s). Refer `ConvolutionParameter.dilation`_. The dilation(s). Refer `ConvolutionParameter.dilation`_.
group : int group : int
The group size. Refer `ConvolutionParameter.group`_. The group size. Refer `ConvolutionParameter.group`_.
weight_filler : FillerParameter weight_filler : FillerParameter
The filler of weights. Refer `ConvolutionParameter.weight_filler`_. The filler of weights. Refer `ConvolutionParameter.weight_filler`_.
bias_filler : FillerParameters bias_filler : FillerParameter
The filler of bias. Refer `ConvolutionParameter.bias_filler`_. The filler of bias. Refer `ConvolutionParameter.bias_filler`_.
""" """
...@@ -76,10 +76,10 @@ class ConvolutionLayer(Layer): ...@@ -76,10 +76,10 @@ class ConvolutionLayer(Layer):
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
inputs = [bottom] + [blob['data'] for blob in self._blobs] inputs = [bottom] + [blob['data'] for blob in self._blobs]
return dragon.ops.Conv2d(inputs, **self.arguments) return _ops.Conv2d(inputs, **self.arguments)
class DepthwiseConvolutionLayer(Layer): class DepthwiseConvolutionLayer(_Layer):
"""The implementation of ``DepthwiseConvolutionLayer``. """The implementation of ``DepthwiseConvolutionLayer``.
Parameters Parameters
...@@ -88,15 +88,15 @@ class DepthwiseConvolutionLayer(Layer): ...@@ -88,15 +88,15 @@ class DepthwiseConvolutionLayer(Layer):
The output channels. Refer `ConvolutionParameter.num_output`_. The output channels. Refer `ConvolutionParameter.num_output`_.
bias_term : boolean bias_term : boolean
Whether to use bias. Refer `ConvolutionParameter.bias_term`_. Whether to use bias. Refer `ConvolutionParameter.bias_term`_.
pad : list of int pad : sequence of int
The zero padding size(s). Refer `ConvolutionParameter.pad`_. The zero padding size(s). Refer `ConvolutionParameter.pad`_.
kernel_size : list of int kernel_size : sequence of int
The kernel size(s). Refer `ConvolutionParameter.kernel_size`_. The kernel size(s). Refer `ConvolutionParameter.kernel_size`_.
stride : list of int stride : sequence of int
The stride(s). Refer `ConvolutionParameter.stride`_. The stride(s). Refer `ConvolutionParameter.stride`_.
weight_filler : FillerParameter weight_filler : FillerParameter
The filler of weights. Refer `ConvolutionParameter.weight_filler`_. The filler of weights. Refer `ConvolutionParameter.weight_filler`_.
bias_filler : FillerParameters bias_filler : FillerParameter
The filler of bias. Refer `ConvolutionParameter.bias_filler`_. The filler of bias. Refer `ConvolutionParameter.bias_filler`_.
""" """
...@@ -130,7 +130,7 @@ class DepthwiseConvolutionLayer(Layer): ...@@ -130,7 +130,7 @@ class DepthwiseConvolutionLayer(Layer):
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
inputs = [bottom] + [blob['data'] for blob in self._blobs] inputs = [bottom] + [blob['data'] for blob in self._blobs]
return dragon.ops.DepthwiseConv2d(inputs, **self.arguments) return _ops.DepthwiseConv2d(inputs, **self.arguments)
class DeconvolutionLayer(ConvolutionLayer): class DeconvolutionLayer(ConvolutionLayer):
...@@ -142,19 +142,19 @@ class DeconvolutionLayer(ConvolutionLayer): ...@@ -142,19 +142,19 @@ class DeconvolutionLayer(ConvolutionLayer):
The output channels. Refer `ConvolutionParameter.num_output`_. The output channels. Refer `ConvolutionParameter.num_output`_.
bias_term : boolean bias_term : boolean
Whether to use bias. Refer `ConvolutionParameter.bias_term`_. Whether to use bias. Refer `ConvolutionParameter.bias_term`_.
pad : list of int pad : sequence of int
The zero padding size(s). Refer `ConvolutionParameter.pad`_. The zero padding size(s). Refer `ConvolutionParameter.pad`_.
kernel_size : list of int kernel_size : sequence of int
The kernel size(s). Refer `ConvolutionParameter.kernel_size`_. The kernel size(s). Refer `ConvolutionParameter.kernel_size`_.
stride : list of int stride : sequence of int
The stride(s). Refer `ConvolutionParameter.stride`_. The stride(s). Refer `ConvolutionParameter.stride`_.
dilation : list of int dilation : sequence of int
The dilation(s). Refer `ConvolutionParameter.dilation`_. The dilation(s). Refer `ConvolutionParameter.dilation`_.
group : int group : int
The group size. Refer `ConvolutionParameter.group`_. The group size. Refer `ConvolutionParameter.group`_.
weight_filler : FillerParameter weight_filler : FillerParameter
The filler of weights. Refer `ConvolutionParameter.weight_filler`_. The filler of weights. Refer `ConvolutionParameter.weight_filler`_.
bias_filler : FillerParameters bias_filler : FillerParameter
The filler of bias. Refer `ConvolutionParameter.bias_filler`_. The filler of bias. Refer `ConvolutionParameter.bias_filler`_.
""" """
...@@ -163,29 +163,29 @@ class DeconvolutionLayer(ConvolutionLayer): ...@@ -163,29 +163,29 @@ class DeconvolutionLayer(ConvolutionLayer):
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
inputs = [bottom] + [blob['data'] for blob in self._blobs] inputs = [bottom] + [blob['data'] for blob in self._blobs]
return dragon.ops.ConvTranspose2d(inputs, **self.arguments) return _ops.ConvTranspose2d(inputs, **self.arguments)
class PoolingLayer(Layer): class PoolingLayer(_Layer):
"""The implementation of ``PoolingLayer``. """The implementation of ``PoolingLayer``.
Parameters Parameters
---------- ----------
pool : PoolMethod pool : PoolMethod
The method. Refer `PoolingParameter.pool`_. The method. Refer `PoolingParameter.pool`_.
pad : list of int pad : sequence of int
The zero padding size(s). Refer `PoolingParameter.pad`_. The zero padding size(s). Refer `PoolingParameter.pad`_.
pad_h : int pad_h : int
The padding size of height. Refer `PoolingParameter.pad_h`_. The padding size of height. Refer `PoolingParameter.pad_h`_.
pad_w : int pad_w : int
The padding size of width. Refer `PoolingParameter.pad_w`_. The padding size of width. Refer `PoolingParameter.pad_w`_.
kernel_size : list of int kernel_size : sequence of int
The kernel size(s). Refer `PoolingParameter.kernel_size`_. The kernel size(s). Refer `PoolingParameter.kernel_size`_.
kernel_h : int kernel_h : int
The kernel size of height. Refer `PoolingParameter.kernel_h`_. The kernel size of height. Refer `PoolingParameter.kernel_h`_.
kernel_w : int kernel_w : int
The kernel size of width. Refer `PoolingParameter.kernel_w`_. The kernel size of width. Refer `PoolingParameter.kernel_w`_.
stride : list of int stride : sequence of int
The strides. Refer `PoolingParameter.stride`_. The strides. Refer `PoolingParameter.stride`_.
stride_h : int stride_h : int
The stride of height. Refer `PoolingParameter.stride_h`_. The stride of height. Refer `PoolingParameter.stride_h`_.
...@@ -212,10 +212,10 @@ class PoolingLayer(Layer): ...@@ -212,10 +212,10 @@ class PoolingLayer(Layer):
else: self.arguments['strides'] = [param.stride_h, param.stride_w] else: self.arguments['strides'] = [param.stride_h, param.stride_w]
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
return dragon.ops.Pool2d(bottom, **self.arguments) return _ops.Pool2d(bottom, **self.arguments)
class ROIPoolingLayer(Layer): class ROIPoolingLayer(_Layer):
"""The implementation of ``ROIPoolingLayer``. """The implementation of ``ROIPoolingLayer``.
Parameters Parameters
...@@ -238,10 +238,10 @@ class ROIPoolingLayer(Layer): ...@@ -238,10 +238,10 @@ class ROIPoolingLayer(Layer):
} }
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
return dragon.ops.ROIPool(bottom, **self.arguments) return _ops.ROIPool(bottom, **self.arguments)
class ROIAlignLayer(Layer): class ROIAlignLayer(_Layer):
"""The implementation of ``ROIAlignLayer``. """The implementation of ``ROIAlignLayer``.
Parameters Parameters
...@@ -264,10 +264,10 @@ class ROIAlignLayer(Layer): ...@@ -264,10 +264,10 @@ class ROIAlignLayer(Layer):
} }
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
return dragon.ops.ROIAlign(bottom, **self.arguments) return _ops.ROIAlign(bottom, **self.arguments)
class LRNLayer(Layer): class LRNLayer(_Layer):
"""The implementation of ``LRNLayer``. """The implementation of ``LRNLayer``.
Parameters Parameters
...@@ -296,15 +296,15 @@ class LRNLayer(Layer): ...@@ -296,15 +296,15 @@ class LRNLayer(Layer):
} }
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
return dragon.ops.LRN(bottom, **self.arguments) return _ops.LRN(bottom, **self.arguments)
class NNResizeLayer(Layer): class NNResizeLayer(_Layer):
"""The implementation of ``NNResizeLayer``. """The implementation of ``NNResizeLayer``.
Parameters Parameters
---------- ----------
shape : caffe_pb2.BlobShape shape : BlobShape
The output shape. Refer `ResizeParameter.shape`_. The output shape. Refer `ResizeParameter.shape`_.
fx : float fx : float
The scale factor of height. Refer `ResizeParameter.fx`_. The scale factor of height. Refer `ResizeParameter.fx`_.
...@@ -330,15 +330,15 @@ class NNResizeLayer(Layer): ...@@ -330,15 +330,15 @@ class NNResizeLayer(Layer):
raise ValueError('The second bottom should be provided to determine the shape.') raise ValueError('The second bottom should be provided to determine the shape.')
self.arguments['shape_like'] = bottom[1] self.arguments['shape_like'] = bottom[1]
bottom = bottom[0] bottom = bottom[0]
return dragon.ops.NNResize(bottom, **self.arguments) return _ops.NNResize(bottom, **self.arguments)
class BilinearResizeLayer(Layer): class BilinearResizeLayer(_Layer):
"""The implementation of ``BilinearResizeLayer``. """The implementation of ``BilinearResizeLayer``.
Parameters Parameters
---------- ----------
shape : caffe_pb2.BlobShape shape : BlobShape
The output shape. Refer `ResizeParameter.shape`_. The output shape. Refer `ResizeParameter.shape`_.
fx : float fx : float
The scale factor of height. Refer `ResizeParameter.fx`_. The scale factor of height. Refer `ResizeParameter.fx`_.
...@@ -364,10 +364,10 @@ class BilinearResizeLayer(Layer): ...@@ -364,10 +364,10 @@ class BilinearResizeLayer(Layer):
raise ValueError('The second bottom should be provided to determine the shape.') raise ValueError('The second bottom should be provided to determine the shape.')
self.arguments['shape_like'] = bottom[1] self.arguments['shape_like'] = bottom[1]
bottom = bottom[0] bottom = bottom[0]
return dragon.ops.BilinearResize(bottom, **self.arguments) return _ops.BilinearResize(bottom, **self.arguments)
class DropBlockLayer(Layer): class DropBlockLayer(_Layer):
"""The implementation of ``DropBlock2dLayer``. """The implementation of ``DropBlock2dLayer``.
Parameters Parameters
...@@ -394,4 +394,4 @@ class DropBlockLayer(Layer): ...@@ -394,4 +394,4 @@ class DropBlockLayer(Layer):
} }
def LayerSetup(self, bottom): def LayerSetup(self, bottom):
return dragon.ops.DropBlock2d(bottom, **self.arguments) return _ops.DropBlock2d(bottom, **self.arguments)
\ No newline at end of file \ No newline at end of file
Dragon/python/dragon/vm/caffe/misc.py
...@@ -15,10 +15,10 @@ from __future__ import absolute_import ...@@ -15,10 +15,10 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import dragon from dragon import config as _cfg
_GLOBAL_ROOT_CAFFE_SOLVER = True _GLOBAL_CAFFE_ROOT_SOLVER = True
def set_mode_cpu(): def set_mode_cpu():
...@@ -33,7 +33,7 @@ def set_mode_cpu(): ...@@ -33,7 +33,7 @@ def set_mode_cpu():
The implementation of `set_mode_cpu(_caffe.cpp, L51)`_. The implementation of `set_mode_cpu(_caffe.cpp, L51)`_.
""" """
dragon.config.EnableCPU() _cfg.EnableCPU()
def set_mode_gpu(): def set_mode_gpu():
...@@ -48,7 +48,7 @@ def set_mode_gpu(): ...@@ -48,7 +48,7 @@ def set_mode_gpu():
The implementation of `set_mode_gpu(_caffe.cpp, L52)`_. The implementation of `set_mode_gpu(_caffe.cpp, L52)`_.
""" """
dragon.config.EnableCUDA() _cfg.EnableCUDA()
def set_device(device): def set_device(device):
...@@ -63,7 +63,7 @@ def set_device(device): ...@@ -63,7 +63,7 @@ def set_device(device):
The implementation of `SetDevice(common.cpp, L65)`_. The implementation of `SetDevice(common.cpp, L65)`_.
""" """
dragon.config.SetGPU(device) _cfg.SetGPU(device)
def set_random_seed(seed): def set_random_seed(seed):
...@@ -83,7 +83,7 @@ def set_random_seed(seed): ...@@ -83,7 +83,7 @@ def set_random_seed(seed):
The implementation of `set_random_seed(_caffe.cpp, L71)`_. The implementation of `set_random_seed(_caffe.cpp, L71)`_.
""" """
dragon.config.SetRandomSeed(seed) _cfg.SetRandomSeed(seed)
def root_solver(): def root_solver():
...@@ -99,7 +99,7 @@ def root_solver(): ...@@ -99,7 +99,7 @@ def root_solver():
The implementation of `root_solver(common.hpp, L164)`_. The implementation of `root_solver(common.hpp, L164)`_.
""" """
return _GLOBAL_ROOT_CAFFE_SOLVER return _GLOBAL_CAFFE_ROOT_SOLVER
def set_root_solver(val): def set_root_solver(val):
...@@ -115,5 +115,5 @@ def set_root_solver(val): ...@@ -115,5 +115,5 @@ def set_root_solver(val):
The implementation of `set_root_solver(common.hpp, L165)`_. The implementation of `set_root_solver(common.hpp, L165)`_.
""" """
global _GLOBAL_ROOT_CAFFE_SOLVER global _GLOBAL_CAFFE_ROOT_SOLVER
_GLOBAL_ROOT_CAFFE_SOLVER = val _GLOBAL_CAFFE_ROOT_SOLVER = val
\ No newline at end of file \ No newline at end of file
Dragon/python/dragon/vm/caffe/net.py
...@@ -15,12 +15,16 @@ from __future__ import absolute_import ...@@ -15,12 +15,16 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import dragon
from collections import OrderedDict from collections import OrderedDict
from google.protobuf.text_format import Parse as parse_text_proto from google.protobuf.text_format import Parse as _parse_text_proto
from dragon.vm.caffe import layers as layer_factory
from dragon.vm.caffe.proto import caffe_pb2 as pb from dragon.core.tensor import Tensor as _Tensor
from dragon.core import workspace as _workspace
from dragon.vm.theano.gradient import grad as _Grad
from dragon.vm.theano.compile.function import function as _Function
from dragon.vm.caffe import layers as _layer_factory
from dragon.vm.caffe.proto import caffe_pb2 as _proto_def
class Blob(object): class Blob(object):
...@@ -89,8 +93,8 @@ class Net(object): ...@@ -89,8 +93,8 @@ class Net(object):
The implementation of `Net_Init(_caffe.cpp, L109)`_. The implementation of `Net_Init(_caffe.cpp, L109)`_.
""" """
self._net = pb.NetParameter() self._net = _proto_def.NetParameter()
parse_text_proto(open(proto_txt,'r').read(), self._net) _parse_text_proto(open(proto_txt,'r').read(), self._net)
self._phase = phase self._phase = phase
self._layers = [] self._layers = []
self._inputs_to_tensors = {} self._inputs_to_tensors = {}
...@@ -100,16 +104,17 @@ class Net(object): ...@@ -100,16 +104,17 @@ class Net(object):
if len(self._net.input) > 0: if len(self._net.input) > 0:
for input in self._net.input: for input in self._net.input:
if not input in self._blobs: if not input in self._blobs:
variable = dragon.Tensor(input).Variable() variable = _Tensor(input).Variable()
self._blobs[input] = { self._blobs[input] = {
'data': variable, 'data': variable,
'diff': dragon.Tensor.Ref(variable.name + '_grad'), 'diff': _Tensor.Ref(variable.name + '_grad'),
} }
self._inputs_to_tensors[input] = self._blobs[input]['data'] self._inputs_to_tensors[input] = self._blobs[input]['data']
for layer in self._net.layer: for layer in self._net.layer:
if not self.FilterLayer(layer): continue if not self.FilterLayer(layer): continue
self._layers.append(getattr(layer_factory, layer.type + 'Layer')(layer)) self._layers.append(getattr(
_layer_factory, layer.type + 'Layer')(layer))
self.Setup() self.Setup()
...@@ -199,7 +204,7 @@ class Net(object): ...@@ -199,7 +204,7 @@ class Net(object):
for idx, top in enumerate(layer._top): for idx, top in enumerate(layer._top):
self._blobs[top] = { self._blobs[top] = {
'data': outputs[idx], 'data': outputs[idx],
'diff': dragon.Tensor.Ref(outputs[idx].name + '_grad'), 'diff': _Tensor.Ref(outputs[idx].name + '_grad'),
} }
self._net_outputs.add(top) self._net_outputs.add(top)
...@@ -271,14 +276,14 @@ class Net(object): ...@@ -271,14 +276,14 @@ class Net(object):
for loss in self.losses: for loss in self.losses:
for var in self.trainable_variables: for var in self.trainable_variables:
dragon.grad(loss, var) _Grad(loss, var)
self._function = dragon.function( self._function = _Function(
outputs=[self.blobs[key].data outputs=[self.blobs[key].data
for key in self.outputs]) for key in self.outputs])
if hasattr(self, '_model'): if hasattr(self, '_model'):
dragon.workspace.Restore(self._model, format='caffe') _workspace.Restore(self._model, format='caffe')
return self._function return self._function
...@@ -299,7 +304,7 @@ class Net(object): ...@@ -299,7 +304,7 @@ class Net(object):
The implementation of `CopyTrainedLayersFromBinaryProto(net.cpp, L780)`_. The implementation of `CopyTrainedLayersFromBinaryProto(net.cpp, L780)`_.
""" """
dragon.workspace.Restore(model, format='caffe') _workspace.Restore(model, format='caffe')
def forward(self, **kwargs): def forward(self, **kwargs):
"""Forward pass. [**PyCaffe Style**] """Forward pass. [**PyCaffe Style**]
...@@ -322,11 +327,11 @@ class Net(object): ...@@ -322,11 +327,11 @@ class Net(object):
def GetOutputs(net, net_outputs): def GetOutputs(net, net_outputs):
ret = {} ret = {}
for output in net_outputs: for output in net_outputs:
ret[output] = dragon.workspace.FetchTensor(net.blobs[output].data) ret[output] = net.blobs[output].data.get_value()
return ret return ret
for name, blob in kwargs.items(): for name, blob in kwargs.items():
dragon.workspace.FeedTensor(self._inputs_to_tensors[name], blob) _workspace.FeedTensor(self._inputs_to_tensors[name], blob)
self.function()(return_outputs=False, stage='forward') self.function()(return_outputs=False, stage='forward')
...@@ -347,7 +352,7 @@ class Net(object): ...@@ -347,7 +352,7 @@ class Net(object):
""" """
for name, blob in kwargs.items(): for name, blob in kwargs.items():
dragon.workspace.FeedTensor(self._inputs_to_tensors[name], blob) _workspace.FeedTensor(self._inputs_to_tensors[name], blob)
self.function()(return_outputs=False, stage='forward') self.function()(return_outputs=False, stage='forward')
def backward(self, **kwargs): def backward(self, **kwargs):
...@@ -368,7 +373,7 @@ class Net(object): ...@@ -368,7 +373,7 @@ class Net(object):
""" """
for name, blob in kwargs.items(): for name, blob in kwargs.items():
dragon.workspace.FeedTensor(self.blobs[name].diff, blob) _workspace.FeedTensor(self.blobs[name].diff, blob)
self.function()(return_outputs=False, stage='backward') self.function()(return_outputs=False, stage='backward')
def save(self, filename): def save(self, filename):
...@@ -399,7 +404,7 @@ class Net(object): ...@@ -399,7 +404,7 @@ class Net(object):
if param.data.name not in keys: if param.data.name not in keys:
tensors.append(param.data) tensors.append(param.data)
keys.add(param.data.name) keys.add(param.data.name)
dragon.workspace.Snapshot(tensors, filename, suffix='', format='caffe') _workspace.Snapshot(tensors, filename, suffix='', format='caffe')
@property @property
def blobs(self): def blobs(self):
......
Dragon/python/dragon/vm/caffe/solver.py
...@@ -16,12 +16,16 @@ from __future__ import division ...@@ -16,12 +16,16 @@ from __future__ import division
from __future__ import print_function from __future__ import print_function
import time import time
import dragon
from google.protobuf.text_format import Parse as parse_text_proto
from dragon.vm.caffe.misc import root_solver from dragon import updaters as _updaters
from dragon.vm.caffe.net import Net from dragon.core import mpi as _mpi
from dragon.vm.caffe.proto import caffe_pb2 as pb from dragon.core import workspace as _workspace
from google.protobuf.text_format import Parse as _parse_text_proto
from dragon.vm.caffe.net import Net as _Net
from dragon.vm.caffe.proto import caffe_pb2 as _proto_def
from dragon.vm.caffe.misc import root_solver as _root_solver
from dragon.vm.theano.compile.function import function as _Function
class Solver(object): class Solver(object):
...@@ -48,8 +52,8 @@ class Solver(object): ...@@ -48,8 +52,8 @@ class Solver(object):
>>> solver = Solver('solver.prototxt') >>> solver = Solver('solver.prototxt')
""" """
self._param = pb.SolverParameter() self._param = _proto_def.SolverParameter()
parse_text_proto(open(proto_txt, 'r').read(), self._param) _parse_text_proto(open(proto_txt, 'r').read(), self._param)
if self._param.iter_size > 1: if self._param.iter_size > 1:
raise NotImplementedError('Gradients accumulating is deprecated.') raise NotImplementedError('Gradients accumulating is deprecated.')
self._net = None self._net = None
...@@ -75,12 +79,12 @@ class Solver(object): ...@@ -75,12 +79,12 @@ class Solver(object):
""" """
if self._param.HasField('net'): if self._param.HasField('net'):
self._net = Net(self._param.net, "TRAIN") self._net = _Net(self._param.net, "TRAIN")
if self._param.HasField('train_net'): if self._param.HasField('train_net'):
if self._net is not None: if self._net is not None:
raise RuntimeError('net or train_net can not be specified both.') raise RuntimeError('net or train_net can not be specified both.')
self._net = Net(self._param.train_net, "TRAIN") self._net = _Net(self._param.train_net, "TRAIN")
def InitTestNets(self): def InitTestNets(self):
"""Initialize the test nets. """Initialize the test nets.
...@@ -94,10 +98,10 @@ class Solver(object): ...@@ -94,10 +98,10 @@ class Solver(object):
The implementation of `InitTestNets(solver.cpp, L104)`_. The implementation of `InitTestNets(solver.cpp, L104)`_.
""" """
if dragon.mpi.Is_Init(): if _mpi.Is_Init():
idx, group = dragon.mpi.AllowParallel() rank, group = _mpi.AllowParallel()
# Only the root in a parallel group can test # Only the root in a parallel group can test
if idx != -1 and dragon.mpi.Rank() != group[0]: return if rank != -1 and _mpi.Rank() != group[0]: return
num_test_net = len(self._param.test_iter) num_test_net = len(self._param.test_iter)
if num_test_net > 0: if num_test_net > 0:
...@@ -106,12 +110,12 @@ class Solver(object): ...@@ -106,12 +110,12 @@ class Solver(object):
if len(self._param.test_net) > 0: if len(self._param.test_net) > 0:
for test_net in self._param.test_net: for test_net in self._param.test_net:
self._test_nets.append(Net(test_net, "TEST")) self._test_nets.append(_Net(test_net, "TEST"))
num_test_net -= len(self._param.test_net) num_test_net -= len(self._param.test_net)
# Consider generic_net # Consider generic_net
if num_test_net > 0: if num_test_net > 0:
self._test_nets.append(Net(self._param.net, "TEST")) self._test_nets.append(_Net(self._param.net, "TEST"))
def BuildNets(self): def BuildNets(self):
"""Build the nets. """Build the nets.
...@@ -164,7 +168,7 @@ class Solver(object): ...@@ -164,7 +168,7 @@ class Solver(object):
blob.decay_multiplier) blob.decay_multiplier)
# Compile # Compile
self.update = dragon.function(updater=self.optimizer) self.update = _Function(updater=self.optimizer)
def GetLearningRate(self): def GetLearningRate(self):
"""Get learning rate based on the preset policy. """Get learning rate based on the preset policy.
...@@ -244,7 +248,7 @@ class Solver(object): ...@@ -244,7 +248,7 @@ class Solver(object):
for iter in range(test_iter): for iter in range(test_iter):
self.tests[test_idx](return_outputs=False) self.tests[test_idx](return_outputs=False)
if not root_solver(): continue if not _root_solver(): continue
if iter == 0: if iter == 0:
for key in net.outputs: for key in net.outputs:
values = net.blobs[key].data.get_value().flatten() values = net.blobs[key].data.get_value().flatten()
...@@ -259,7 +263,7 @@ class Solver(object): ...@@ -259,7 +263,7 @@ class Solver(object):
test_score[i] += value test_score[i] += value
i += 1 i += 1
if not root_solver(): return if not _root_solver(): return
print('Iteration {}, Test net #{}'.format(self.iter, test_idx)) print('Iteration {}, Test net #{}'.format(self.iter, test_idx))
for idx, score in enumerate(test_score): for idx, score in enumerate(test_score):
...@@ -299,12 +303,12 @@ class Solver(object): ...@@ -299,12 +303,12 @@ class Solver(object):
loss = 0.0 loss = 0.0
for i in range(self._param.iter_size): for i in range(self._param.iter_size):
self.train(return_outputs=False) self.train(return_outputs=False)
if root_solver(): if _root_solver():
for e in self.net.losses: for e in self.net.losses:
values = e.get_value().flatten() values = e.get_value().flatten()
for v in values: loss += v for v in values: loss += v
if root_solver(): if _root_solver():
loss /= self._param.iter_size loss /= self._param.iter_size
if len(loss_vec) < self._param.average_loss: if len(loss_vec) < self._param.average_loss:
loss_vec.append(loss) loss_vec.append(loss)
...@@ -319,7 +323,7 @@ class Solver(object): ...@@ -319,7 +323,7 @@ class Solver(object):
self.update() self.update()
# Display # Display
if root_solver() and self._param.display: if _root_solver() and self._param.display:
if self.iter % self._param.display == 0: if self.iter % self._param.display == 0:
base_lr = self.optimizer.base_lr base_lr = self.optimizer.base_lr
print('Iteration %d, lr = %s, loss = %f, time = %.2fs' % \ print('Iteration %d, lr = %s, loss = %f, time = %.2fs' % \
...@@ -410,7 +414,7 @@ class Solver(object): ...@@ -410,7 +414,7 @@ class Solver(object):
""" """
tensors = [blob.data for blob in self._layer_blobs] tensors = [blob.data for blob in self._layer_blobs]
filename = "_iter_" + str(self.iter) filename = "_iter_" + str(self.iter)
dragon.workspace.Snapshot(tensors, filename, _workspace.Snapshot(tensors, filename,
prefix=self._param.snapshot_prefix, prefix=self._param.snapshot_prefix,
suffix='.caffemodel', format='caffe') suffix='.caffemodel', format='caffe')
...@@ -492,7 +496,7 @@ class SGDSolver(Solver): ...@@ -492,7 +496,7 @@ class SGDSolver(Solver):
""" """
def __init__(self, proto_txt): def __init__(self, proto_txt):
super(SGDSolver, self).__init__(proto_txt=proto_txt) super(SGDSolver, self).__init__(proto_txt=proto_txt)
self.optimizer = dragon.updaters.SGDUpdater(**self._optimizer_arguments) self.optimizer = _updaters.SGDUpdater(**self._optimizer_arguments)
self.BuildOptimizer() self.BuildOptimizer()
def ParseOptimizerArguments(self): def ParseOptimizerArguments(self):
...@@ -514,7 +518,7 @@ class NesterovSolver(Solver): ...@@ -514,7 +518,7 @@ class NesterovSolver(Solver):
""" """
def __init__(self, proto_txt): def __init__(self, proto_txt):
super(NesterovSolver, self).__init__(proto_txt=proto_txt) super(NesterovSolver, self).__init__(proto_txt=proto_txt)
self.optimizer = dragon.updaters.NesterovUpdater(**self._optimizer_arguments) self.optimizer = _updaters.NesterovUpdater(**self._optimizer_arguments)
self.BuildOptimizer() self.BuildOptimizer()
def ParseOptimizerArguments(self): def ParseOptimizerArguments(self):
...@@ -538,7 +542,7 @@ class RMSPropSolver(Solver): ...@@ -538,7 +542,7 @@ class RMSPropSolver(Solver):
""" """
def __init__(self, proto_txt): def __init__(self, proto_txt):
super(RMSPropSolver, self).__init__(proto_txt=proto_txt) super(RMSPropSolver, self).__init__(proto_txt=proto_txt)
self.optimizer = dragon.updaters.RMSPropUpdater(**self._optimizer_arguments) self.optimizer = _updaters.RMSPropUpdater(**self._optimizer_arguments)
self.BuildOptimizer() self.BuildOptimizer()
def ParseOptimizerArguments(self): def ParseOptimizerArguments(self):
...@@ -565,7 +569,7 @@ class AdamSolver(Solver): ...@@ -565,7 +569,7 @@ class AdamSolver(Solver):
""" """
def __init__(self, proto_txt): def __init__(self, proto_txt):
super(AdamSolver, self).__init__(proto_txt=proto_txt) super(AdamSolver, self).__init__(proto_txt=proto_txt)
self.optimizer = dragon.updaters.AdamUpdater(**self._optimizer_arguments) self.optimizer = _updaters.AdamUpdater(**self._optimizer_arguments)
self.BuildOptimizer() self.BuildOptimizer()
def ParseOptimizerArguments(self): def ParseOptimizerArguments(self):
......
Dragon/python/dragon/vm/onnx/frontend.py
...@@ -17,17 +17,18 @@ from __future__ import absolute_import ...@@ -17,17 +17,18 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import numpy
import itertools import itertools
import numpy as np
from collections import defaultdict from collections import defaultdict
from onnx import (checker, mapping, numpy_helper, GraphProto, OperatorSetIdProto) from onnx import checker, mapping, numpy_helper, GraphProto, OperatorSetIdProto
from onnx.helper import make_tensor_value_info, make_model, printable_graph from onnx.helper import make_tensor_value_info, make_model, printable_graph
from dragon.vm.onnx.helper import \ from dragon.core import workspace as _workspace
(extract_initializer, extract_leaf_tensors, from dragon.vm.onnx.helper import native_run_graph
native_run_graph, fetch_initializer,) from dragon.vm.onnx.helper import fetch_initializer
from dragon.vm.onnx.helper import extract_initializer
from dragon.vm.onnx.helper import extract_leaf_tensors
from dragon.vm.onnx.nodes.factory import get_nodes_def from dragon.vm.onnx.nodes.factory import get_nodes_def
...@@ -104,15 +105,22 @@ class DragonFrontend(object): ...@@ -104,15 +105,22 @@ class DragonFrontend(object):
if run_native_graph and not enforce_no_running: if run_native_graph and not enforce_no_running:
inputs = {} inputs = {}
for name, (elem_type, shape) in value_info.items(): for name, (elem_type, shape) in value_info.items():
inputs[name] = np.random.randn(*shape).astype( inputs[name] = numpy.random.randn(*shape).astype(
mapping.TENSOR_TYPE_TO_NP_TYPE[elem_type]) mapping.TENSOR_TYPE_TO_NP_TYPE[elem_type])
ws, outputs, initializer = native_run_graph( ws, outputs, initializer = native_run_graph(
graph_def, inputs, initializer, init_func) graph_def, inputs, initializer, init_func)
for name in graph_def.output:
output = outputs[name]
elem_type = mapping.NP_TYPE_TO_TENSOR_TYPE[output.dtype]
shape = output.shape
value_info[name] = (elem_type, shape)
if enforce_no_running: if enforce_no_running:
# In some cases(e.g. PyTorch), we had ran the graph # In some cases(e.g. PyTorch), we had ran the graph
# outputs had been in ``value_info`` already # outputs had been in ``value_info`` already
import dragon.core.workspace as ws ws = _workspace.get_default_workspace()
initializer = fetch_initializer(initializer) initializer = fetch_initializer(initializer)
# Prepare to make the graph # Prepare to make the graph
......
Dragon/python/dragon/vm/onnx/helper.py
...@@ -21,8 +21,8 @@ import sys ...@@ -21,8 +21,8 @@ import sys
from onnx.backend.base import namedtupledict from onnx.backend.base import namedtupledict
from onnx import numpy_helper from onnx import numpy_helper
import dragon as dg from dragon.core import workspace as _workspace
from dragon.vm.onnx.workspace import Workspace from dragon.core.tensor import Tensor as _Tensor
INITIALIZER_TAG = { INITIALIZER_TAG = {
...@@ -65,7 +65,7 @@ def fetch_initializer(initializer): ...@@ -65,7 +65,7 @@ def fetch_initializer(initializer):
# Fetch the initializer # Fetch the initializer
return [ return [
numpy_helper.from_array( numpy_helper.from_array(
dg.workspace.FetchTensor(name), name=name) _workspace.FetchTensor(name), name=name)
for name in initializer for name in initializer
] ]
...@@ -87,32 +87,32 @@ def native_run_graph(graph_def, inputs, initializer, init_func=None): ...@@ -87,32 +87,32 @@ def native_run_graph(graph_def, inputs, initializer, init_func=None):
graph_def.arg[i].i = 0 graph_def.arg[i].i = 0
# Create an anonymous workspace # Create an anonymous workspace
ws = Workspace() ws = _workspace.Workspace()
with dg.ws_scope(ws.name): with ws.as_default():
# Register all the initializer before feeding them # Register all the initializer before feeding them
for name in initializer: for name in initializer:
dg.Tensor(name=name).Variable() _Tensor(name=name).Variable()
# Feed the given values if necessary # Feed the given values if necessary
if init_func: init_func() if init_func: init_func()
# Feed the external inputs # Feed the external inputs
for name, blob in inputs.items(): for name, blob in inputs.items():
dg.workspace.FeedTensor(name, blob) _workspace.FeedTensor(name, blob)
# Create and Run the graph # Create and Run the graph
graph_name = dg.workspace.CreateGraph(graph_def) graph_name = _workspace.CreateGraph(graph_def)
dg.workspace.RunGraph(graph_name, return_outputs=False) _workspace.RunGraph(graph_name, return_outputs=False)
# Fetch the outputs # Fetch the outputs
output_names = graph_def.output output_names = graph_def.output
output_values = [dg.workspace.FetchTensor(name) for name in output_names] output_values = [_workspace.FetchTensor(name) for name in output_names]
# Fetch the initializer # Fetch the initializer
initializer = [ initializer = [
numpy_helper.from_array( numpy_helper.from_array(
dg.workspace.FetchTensor(name), name=name) _workspace.FetchTensor(name), name=name)
for name in initializer for name in initializer
] ]
......
Dragon/python/dragon/vm/onnx/utils.py
...@@ -16,12 +16,12 @@ from __future__ import division ...@@ -16,12 +16,12 @@ from __future__ import division
from __future__ import print_function from __future__ import print_function
import os import os
import numpy as np import numpy
from onnx import mapping
from google.protobuf.text_format import Parse as parse_text_proto
import dragon.proto.dragon_pb2 as pb from onnx import mapping as _mapping
import dragon.import_c_api as C from dragon.core import workspace as _workspace
from dragon.proto import dragon_pb2 as _proto_def
from google.protobuf.text_format import Parse as _parse_text_proto
from dragon.vm.theano.compile.function import Function from dragon.vm.theano.compile.function import Function
from dragon.vm.onnx.frontend import graph_def_to_onnx_model from dragon.vm.onnx.frontend import graph_def_to_onnx_model
...@@ -119,8 +119,8 @@ def export_from_graph_text( ...@@ -119,8 +119,8 @@ def export_from_graph_text(
""" """
with open(text_file, 'r') as rf: with open(text_file, 'r') as rf:
graph_def = pb.GraphDef() graph_def = _proto_def.GraphDef()
parse_text_proto(rf.read(), graph_def) _parse_text_proto(rf.read(), graph_def)
export_from_graph_def( export_from_graph_def(
graph_def=graph_def, graph_def=graph_def,
...@@ -148,8 +148,10 @@ def import_to_graph_def(model_path): ...@@ -148,8 +148,10 @@ def import_to_graph_def(model_path):
""" """
if not os.path.exists(model_path): if not os.path.exists(model_path):
raise ValueError('Given model({}) is not existed.'.format(model_path)) raise ValueError('Given model({}) is not existed.'.format(model_path))
graph_def = pb.GraphDef() graph_def = _proto_def.GraphDef()
serialized_proto = C.ImportONNXModel(model_path) serialized_proto = _workspace \
.get_default_workspace() \
.ImportONNXModel(model_path)
graph_def.ParseFromString(serialized_proto) graph_def.ParseFromString(serialized_proto)
return graph_def return graph_def
...@@ -238,4 +240,4 @@ def surgery_on_graph_def( ...@@ -238,4 +240,4 @@ def surgery_on_graph_def(
def make_value_info(shape, dtype='float32'): def make_value_info(shape, dtype='float32'):
return mapping.NP_TYPE_TO_TENSOR_TYPE[np.dtype(dtype)], shape return _mapping.NP_TYPE_TO_TENSOR_TYPE[numpy.dtype(dtype)], shape
\ No newline at end of file \ No newline at end of file
Dragon/python/dragon/vm/onnx/workspace.py deleted 100644 → 0
# ------------------------------------------------------------
# Copyright (c) 2017-present, SeetaTech, Co.,Ltd.
#
# Licensed under the BSD 2-Clause License.
# You should have received a copy of the BSD 2-Clause License
# along with the software. If not, See,
#
# <https://opensource.org/licenses/BSD-2-Clause>
#
# Codes are based on:
#
# <https://github.com/pytorch/pytorch/blob/master/caffe2/python/onnx/workspace.py>
#
# ------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import uuid
import dragon as dg
class Workspace(object):
def __init__(self):
self.name = 'onnx/' + str(uuid.uuid4())
def __getattr__(self, attr):
def f(*args, **kwargs):
with dg.ws_scope(self.name, ):
return getattr(dg.workspace, attr)(*args, **kwargs)
return f
def __del__(self):
self.ResetWorkspace(self.name)
\ No newline at end of file
Dragon/python/dragon/vm/tensorflow/client/session.py
...@@ -16,45 +16,42 @@ from __future__ import print_function ...@@ -16,45 +16,42 @@ from __future__ import print_function
import warnings import warnings
from collections import defaultdict from collections import defaultdict
import dragon from dragon.core import workspace as _workspace
from dragon.core.tensor import Tensor as _Tensor
from dragon.vm.theano.compile import function as _Function
from dragon.vm.tensorflow.protobuf import config_pb2 from dragon.vm.tensorflow.protobuf import config_pb2
from dragon.vm.tensorflow.training.optimizer import Optimizer from dragon.vm.tensorflow.training.optimizer import Optimizer
from dragon.vm.tensorflow.ops.variables import VariablesInitializer from dragon.vm.tensorflow.ops.variables import VariablesInitializer
from dragon.vm.tensorflow.framework import ops from dragon.vm.tensorflow.framework import ops
_GLOBAL_DATA_FLOW_KEYS = defaultdict(dict)
class _DataFlow(object): class _DataFlow(object):
"""DataFlow takes a group of expressions and """DataFlow takes a group of expressions and
the specified output tensors. the specified output tensors.
We store the flows that requiring the same output names, We store the flows that requiring the same output names,
i.e., those flows can be reused and should not to create a new graph. i.e., those flows can be reused and should not be created again.
""" """
def __init__(self, functions): def __init__(self, functions):
self.functions = functions self.functions = functions
def run(self, feed_dict=None): def run(self, feed_dict=None):
for i, function in enumerate(self.functions): for i, func in enumerate(self.functions):
if i == 0 and feed_dict is not None: if i == 0 and feed_dict is not None:
for tensor, value in feed_dict.items(): for tensor, value in feed_dict.items():
dragon.workspace.FeedTensor(tensor, value) _workspace.FeedTensor(tensor, value)
function(return_outputs=False) func(return_outputs=False)
@classmethod @classmethod
def try_get(cls, workspace, flow_key): def try_get(cls, graph_id, flow_key):
global _GLOBAL_DATA_FLOW_KEYS if flow_key in _GLOBAL_DATA_FLOWS[graph_id]:
if flow_key in _GLOBAL_DATA_FLOW_KEYS[workspace]: return _GLOBAL_DATA_FLOWS[graph_id][flow_key]
return _GLOBAL_DATA_FLOW_KEYS[workspace][flow_key]
@classmethod @classmethod
def try_add(cls, workspace, flow_key, flow): def try_add(cls, graph_id, flow_key, flow):
global _GLOBAL_DATA_FLOW_KEYS global _GLOBAL_DATA_FLOWS
_GLOBAL_DATA_FLOW_KEYS[workspace][flow_key] = flow _GLOBAL_DATA_FLOWS[graph_id][flow_key] = flow
class BaseSession(object): class BaseSession(object):
...@@ -115,7 +112,7 @@ class BaseSession(object): ...@@ -115,7 +112,7 @@ class BaseSession(object):
for e in fetches: for e in fetches:
if isinstance(e, Optimizer): optimizers.append(e) if isinstance(e, Optimizer): optimizers.append(e)
elif isinstance(e, VariablesInitializer): tensors.extend(e.var_list) elif isinstance(e, VariablesInitializer): tensors.extend(e.var_list)
elif isinstance(e, dragon.Tensor): tensors.append(e) elif isinstance(e, _Tensor): tensors.append(e)
# Find minimum solving targets # Find minimum solving targets
targets = set() targets = set()
...@@ -124,45 +121,45 @@ class BaseSession(object): ...@@ -124,45 +121,45 @@ class BaseSession(object):
for t in optimizer._targets: targets.add(t) for t in optimizer._targets: targets.add(t)
targets = list(targets) targets = list(targets)
gen_flow_key = tuple(e.name for e in targets) flow_key = tuple(e.name for e in targets)
# Exist this data flow before? # Exist this data flow before?
data_flow = _DataFlow.try_get( flow = _DataFlow.try_get(id(self._graph), flow_key)
self._graph._workspace, gen_flow_key)
# Run by feeding # Run by feeding
if feed_dict is not None: if feed_dict is not None:
# Check the feed dict # Check the feed dict
for key, value in feed_dict.items(): for key, value in feed_dict.items():
if not isinstance(key, dragon.Tensor): if not isinstance(key, _Tensor):
raise TypeError('The key of feed_dict key should be a Tensor.') raise TypeError('The key of ``feed_dict`` should be a Tensor.')
if key.shape is not None: if key.shape is not None:
# Align the number of dimensions # Align the number of dimensions
if len(key.shape) != len(value.shape): if len(key.shape) != len(value.shape):
raise RuntimeError( raise RuntimeError(
'The Tensor({}) was limited to {} dimensions, \ 'The Tensor({}) was limited to {} dimensions, '\
while feed a value with {} dimensions.' 'while feed a value with {} dimensions.'
.format(key.name, len(key.shape), len(value.shape))) .format(key.name, len(key.shape), len(value.shape)))
# Verify for the each dimension # Verify for the each dimension
for i in range(len(key.shape)): for i in range(len(key.shape)):
if key.shape[i] is None: continue if key.shape[i] is None: continue
if key.shape[i] != value.shape[i]: if key.shape[i] != value.shape[i]:
raise RuntimeError( raise RuntimeError(
'The shape of Tensor({}) was limited as ('.format(key.name) + 'The shape of Tensor({}) was limited as ('.format(key.name) +
','.join([str(dim) for dim in key.shape]) + '), ' + ','.join([str(dim) for dim in key.shape]) + '), ' +
'while feed a value with (' + ','.join([str(dim) for dim in value.shape]) + ').') 'while feed a value with (' +
','.join([str(dim) for dim in value.shape]) + ').')
# Create a new data flow if necessary # Create a new data flow if necessary
if data_flow is None: if flow is None:
functions = [dragon.function(outputs=targets)] functions = [_Function(outputs=targets)]
for optimizer in optimizers: for optimizer in optimizers:
functions.append(dragon.function( functions.append(_Function(
updater=optimizer.updater)) updater=optimizer.updater))
data_flow = _DataFlow(functions) flow = _DataFlow(functions)
_DataFlow.try_add(self.graph._workspace, gen_flow_key, data_flow) _DataFlow.try_add(id(self._graph), flow_key, flow)
# Run this data flow # Run this data flow
data_flow.run(feed_dict) flow.run(feed_dict)
# Fetch after running # Fetch after running
returns = [] returns = []
...@@ -233,4 +230,9 @@ class InteractiveSession(BaseSession): ...@@ -233,4 +230,9 @@ class InteractiveSession(BaseSession):
@staticmethod @staticmethod
def reset(target, containers=None, config=None): def reset(target, containers=None, config=None):
pass pass
\ No newline at end of file
# Store the flows for different graphs
# ThreadLocal is not necessary
_GLOBAL_DATA_FLOWS = defaultdict(dict)
\ No newline at end of file
Dragon/python/dragon/vm/tensorflow/contrib/framework/ops/variables.py
...@@ -13,8 +13,11 @@ from dragon.vm.tensorflow.framework import ops ...@@ -13,8 +13,11 @@ from dragon.vm.tensorflow.framework import ops
from dragon.vm.tensorflow.ops import var_scope as variable_scope from dragon.vm.tensorflow.ops import var_scope as variable_scope
def get_variables(scope=None, suffix=None, def get_variables(
collection=ops.GraphKeys.GLOBAL_VARIABLES): scope=None,
suffix=None,
collection=ops.GraphKeys.GLOBAL_VARIABLES,
):
if isinstance(scope, variable_scope.VariableScope): if isinstance(scope, variable_scope.VariableScope):
scope = scope.name scope = scope.name
if suffix is not None: if suffix is not None:
......
Dragon/python/dragon/vm/tensorflow/contrib/layers/initializers.py
...@@ -19,25 +19,31 @@ from __future__ import print_function ...@@ -19,25 +19,31 @@ from __future__ import print_function
import math import math
from dragon.vm.tensorflow.framework import dtypes
from dragon.vm.tensorflow.ops import random_ops from dragon.vm.tensorflow.ops import random_ops
from dragon.vm.tensorflow.framework import dtypes
__all__ = ['xavier_initializer', def xavier_initializer(
'xavier_initializer_conv2d', uniform=True,
'variance_scaling_initializer'] seed=None,
dtype=dtypes.float32,
):
def xavier_initializer(uniform=True, seed=None, dtype=dtypes.float32): return variance_scaling_initializer(
return variance_scaling_initializer(factor=1.0, mode='FAN_AVG', factor=1.0,
uniform=uniform, seed=seed, dtype=dtype) mode='FAN_AVG',
uniform=uniform,
seed=seed,
xavier_initializer_conv2d = xavier_initializer dtype=dtype,
)
def variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, def variance_scaling_initializer(
seed=None, dtype=dtypes.float32): factor=2.0,
mode='FAN_IN',
uniform=False,
seed=None,
dtype=dtypes.float32,
):
if not dtype.is_floating: if not dtype.is_floating:
raise TypeError('Cannot create initializer for non-floating point type.') raise TypeError('Cannot create initializer for non-floating point type.')
if mode not in ['FAN_IN', 'FAN_OUT', 'FAN_AVG']: if mode not in ['FAN_IN', 'FAN_OUT', 'FAN_AVG']:
...@@ -79,3 +85,7 @@ def variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, ...@@ -79,3 +85,7 @@ def variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False,
seed=seed) seed=seed)
return _initializer return _initializer
# Alias
xavier_initializer_conv2d = xavier_initializer
\ No newline at end of file
Dragon/python/dragon/vm/tensorflow/contrib/layers/layers.py
...@@ -17,20 +17,14 @@ from __future__ import absolute_import ...@@ -17,20 +17,14 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
from collections import defaultdict from dragon.vm.tensorflow.framework import ops
import dragon.ops as op_lib
import dragon.vm.tensorflow.framework.ops as ops
from dragon.vm.tensorflow.contrib.layers import initializers from dragon.vm.tensorflow.contrib.layers import initializers
from dragon.vm.tensorflow.ops import init_ops from dragon.vm.tensorflow.ops import init_ops
from dragon.vm.tensorflow.ops import nn from dragon.vm.tensorflow.ops import nn
from dragon.vm.tensorflow.ops import var_scope as vs from dragon.vm.tensorflow.ops import var_scope as vs
from dragon.vm.tensorflow.layers import layers from dragon.vm.tensorflow.layers import layers
from dragon.ops import Flatten as _FlattenOp
__all__ = ['flatten']
_LAYERS_UID_DICT = defaultdict(int)
DATA_FORMAT_NCHW = 'NCHW' DATA_FORMAT_NCHW = 'NCHW'
DATA_FORMAT_NHWC = 'NHWC' DATA_FORMAT_NHWC = 'NHWC'
...@@ -38,72 +32,72 @@ DATA_FORMAT_NCDHW = 'NCDHW' ...@@ -38,72 +32,72 @@ DATA_FORMAT_NCDHW = 'NCDHW'
DATA_FORMAT_NDHWC = 'NDHWC' DATA_FORMAT_NDHWC = 'NDHWC'
def _default_scope(scope, key, indicator): def avg_pool2d(
if scope is None: inputs,
return indicator kernel_size,
# global _LAYERS_UID_DICT stride=2,
# _LAYERS_UID_DICT[key] += 1 padding='VALID',
# return '{}{}'.format(indicator, _LAYERS_UID_DICT[key]) data_format=DATA_FORMAT_NHWC,
else: outputs_collections=None,
return scope scope=None,
):
def avg_pool2d(inputs,
kernel_size,
stride=2,
padding='VALID',
data_format=DATA_FORMAT_NHWC,
outputs_collections=None,
scope=None):
if data_format not in (DATA_FORMAT_NCHW, DATA_FORMAT_NHWC): if data_format not in (DATA_FORMAT_NCHW, DATA_FORMAT_NHWC):
raise ValueError('data_format has to be either NCHW or NHWC.') raise ValueError('data_format has to be either NCHW or NHWC.')
df = ('channels_first' if data_format and data_format.startswith('NC') df = ('channels_first' if data_format and data_format.startswith('NC')
else 'channels_last') else 'channels_last')
return layers.average_pooling2d(inputs=inputs, return layers.average_pooling2d(
pool_size=kernel_size, inputs=inputs,
strides=stride, pool_size=kernel_size,
padding=padding, strides=stride,
data_format=df) padding=padding,
data_format=df,
)
def max_pool2d(inputs,
kernel_size,
stride=2, def max_pool2d(
padding='VALID', inputs,
data_format=DATA_FORMAT_NHWC, kernel_size,
outputs_collections=None, stride=2,
scope=None): padding='VALID',
data_format=DATA_FORMAT_NHWC,
outputs_collections=None,
scope=None,
):
if data_format not in (DATA_FORMAT_NCHW, DATA_FORMAT_NHWC): if data_format not in (DATA_FORMAT_NCHW, DATA_FORMAT_NHWC):
raise ValueError('data_format has to be either NCHW or NHWC.') raise ValueError('data_format has to be either NCHW or NHWC.')
df = ('channels_first' if data_format and data_format.startswith('NC') df = ('channels_first' if data_format and data_format.startswith('NC')
else 'channels_last') else 'channels_last')
return layers.max_pooling2d(inputs=inputs, return layers.max_pooling2d(
pool_size=kernel_size, inputs=inputs,
strides=stride, pool_size=kernel_size,
padding=padding, strides=stride,
data_format=df) padding=padding,
data_format=df,
)
def convolution(inputs,
num_outputs,
kernel_size, def convolution(
stride=1, inputs,
padding='SAME', num_outputs,
data_format=None, kernel_size,
rate=1, stride=1,
activation_fn=nn.relu, padding='SAME',
normalizer_fn=None, data_format=None,
normalizer_params=None, rate=1,
weights_initializer=initializers.xavier_initializer(), activation_fn=nn.relu,
weights_regularizer=None, normalizer_fn=None,
biases_initializer=init_ops.zeros_initializer(), normalizer_params=None,
biases_regularizer=None, weights_initializer=initializers.xavier_initializer(),
reuse=None, weights_regularizer=None,
variables_collections=None, biases_initializer=init_ops.zeros_initializer(),
outputs_collections=None, biases_regularizer=None,
trainable=True, reuse=None,
scope=None): variables_collections=None,
scope = _default_scope(scope, 'CONVOLUTION', 'Conv') outputs_collections=None,
trainable=True,
scope=None,
):
scope = _default_scope(scope, 'Conv')
if data_format not in [None, 'NHWC', 'NCHW']: if data_format not in [None, 'NHWC', 'NCHW']:
raise ValueError('Invalid data_format: %r' % (data_format,)) raise ValueError('Invalid data_format: %r' % (data_format,))
data_format = 'channels_first' if data_format == 'NCHW' else 'channels_last' data_format = 'channels_first' if data_format == 'NCHW' else 'channels_last'
...@@ -126,7 +120,8 @@ def convolution(inputs, ...@@ -126,7 +120,8 @@ def convolution(inputs,
bias_regularizer=biases_regularizer, bias_regularizer=biases_regularizer,
activity_regularizer=None, activity_regularizer=None,
trainable=trainable, trainable=trainable,
reuse=reuse) reuse=reuse,
)
# Simple alias. # Simple alias.
...@@ -134,21 +129,23 @@ convolution2d = convolution ...@@ -134,21 +129,23 @@ convolution2d = convolution
conv2d = convolution2d conv2d = convolution2d
def fully_connected(inputs, def fully_connected(
num_outputs, inputs,
activation_fn=nn.relu, num_outputs,
normalizer_fn=None, activation_fn=nn.relu,
normalizer_params=None, normalizer_fn=None,
weights_initializer=initializers.xavier_initializer(), normalizer_params=None,
weights_regularizer=None, weights_initializer=initializers.xavier_initializer(),
biases_initializer=init_ops.zeros_initializer(), weights_regularizer=None,
biases_regularizer=None, biases_initializer=init_ops.zeros_initializer(),
reuse=None, biases_regularizer=None,
variables_collections=None, reuse=None,
outputs_collections=None, variables_collections=None,
trainable=True, outputs_collections=None,
scope=None): trainable=True,
scope = _default_scope(scope, 'FULLY_CONNECTED', 'fully_connected') scope=None,
):
scope = _default_scope(scope, 'fully_connected')
with vs.variable_scope(scope, reuse=reuse) as sc: with vs.variable_scope(scope, reuse=reuse) as sc:
return layers.dense( return layers.dense(
inputs=inputs, inputs=inputs,
...@@ -160,32 +157,35 @@ def fully_connected(inputs, ...@@ -160,32 +157,35 @@ def fully_connected(inputs,
bias_regularizer=biases_regularizer, bias_regularizer=biases_regularizer,
activity_regularizer=None, activity_regularizer=None,
trainable=trainable, trainable=trainable,
reuse=reuse) reuse=reuse,
)
def batch_norm(inputs,
decay=0.999, def batch_norm(
center=True, inputs,
scale=False, decay=0.999,
epsilon=0.001, center=True,
activation_fn=None, scale=False,
param_initializers=None, epsilon=0.001,
param_regularizers=None, activation_fn=None,
updates_collections=ops.GraphKeys.UPDATE_OPS, param_initializers=None,
is_training=True, param_regularizers=None,
reuse=None, updates_collections=ops.GraphKeys.UPDATE_OPS,
variables_collections=None, is_training=True,
outputs_collections=None, reuse=None,
trainable=True, variables_collections=None,
batch_weights=None, outputs_collections=None,
fused=False, trainable=True,
data_format=DATA_FORMAT_NHWC, batch_weights=None,
zero_debias_moving_mean=False, fused=False,
scope=None, data_format=DATA_FORMAT_NHWC,
renorm=False, zero_debias_moving_mean=False,
renorm_clipping=None, scope=None,
renorm_decay=0.99): renorm=False,
scope = _default_scope(scope, 'BATCH_NORM', 'BatchNorm') renorm_clipping=None,
renorm_decay=0.99,
):
scope = _default_scope(scope, 'BatchNorm')
if data_format not in (DATA_FORMAT_NCHW, DATA_FORMAT_NHWC): if data_format not in (DATA_FORMAT_NCHW, DATA_FORMAT_NHWC):
raise ValueError('data_format has to be either NCHW or NHWC.') raise ValueError('data_format has to be either NCHW or NHWC.')
axis = 1 if data_format == DATA_FORMAT_NCHW else -1 axis = 1 if data_format == DATA_FORMAT_NCHW else -1
...@@ -193,10 +193,14 @@ def batch_norm(inputs, ...@@ -193,10 +193,14 @@ def batch_norm(inputs,
with vs.variable_scope(scope, reuse=reuse) as sc: with vs.variable_scope(scope, reuse=reuse) as sc:
if not param_initializers: if not param_initializers:
param_initializers = {} param_initializers = {}
beta_initializer = param_initializers.get('beta', init_ops.zeros_initializer()) beta_initializer = param_initializers.get(
gamma_initializer = param_initializers.get('gamma', init_ops.ones_initializer()) 'beta', init_ops.zeros_initializer())
moving_mean_initializer = param_initializers.get('moving_mean', init_ops.zeros_initializer()) gamma_initializer = param_initializers.get(
moving_variance_initializer = param_initializers.get('moving_variance', init_ops.ones_initializer()) 'gamma', init_ops.ones_initializer())
moving_mean_initializer = param_initializers.get(
'moving_mean', init_ops.zeros_initializer())
moving_variance_initializer = param_initializers.get(
'moving_variance', init_ops.ones_initializer())
if not param_regularizers: if not param_regularizers:
param_regularizers = {} param_regularizers = {}
...@@ -222,11 +226,19 @@ def batch_norm(inputs, ...@@ -222,11 +226,19 @@ def batch_norm(inputs,
renorm_clipping=renorm_clipping, renorm_clipping=renorm_clipping,
renorm_momentum=renorm_decay, renorm_momentum=renorm_decay,
fused=fused, fused=fused,
training=is_training) training=is_training,
)
def flatten(
inputs,
outputs_collections=None,
scope=None,
):
return _FlattenOp(inputs, axis=0, keep_axes=2)
def flatten(inputs,
outputs_collections=None,
scope=None):
return op_lib.Flatten(inputs, axis=0, keep_axes=2)
def _default_scope(scope, indicator):
"""Return the default scope."""
if scope is None: return indicator
else: return scope
\ No newline at end of file
Dragon/python/dragon/vm/tensorflow/framework/constant_op.py
...@@ -13,60 +13,68 @@ from __future__ import absolute_import ...@@ -13,60 +13,68 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import dragon import numpy
import numpy as np
from dragon.core import scope as _scope
from dragon.core import workspace as _workspace
from dragon.core.tensor import Tensor as _Tensor
def constant(value, dtype=None, shape=None, name=None, verify_shape=False):
def constant(
value,
dtype=None,
shape=None,
name=None,
verify_shape=False,
):
if dtype is not None: if dtype is not None:
if isinstance(value, np.ndarray): if isinstance(value, numpy.ndarray):
feed = value.astype(dtype.as_numpy_dtype) value = value.astype(dtype.as_numpy_dtype)
elif isinstance(value, list):
feed = np.array(value, dtype.as_numpy_dtype)
else: else:
feed = np.array([value], dtype.as_numpy_dtype) value = numpy.array(value, dtype.as_numpy_dtype)
else: else:
if isinstance(value, np.ndarray): feed = value if not isinstance(value, numpy.ndarray):
else: value = numpy.array(value)
feed = np.array(value)
# Discard the default float64 # Discard the default float64
if feed.dtype == np.float64: if value.dtype == numpy.float64:
feed = feed.astype(np.float32) value = value.astype(numpy.float32)
# Determine the shape # Determine the shape
if shape is not None: if shape is not None:
if feed.size == 1: if value.size == 1:
# Case 1: Broadcast with scalar value # Case 1: Broadcast with scalar value
c = feed.flatten()[0] scalar = value.flatten()[0]
feed = np.zeros(shape, feed.dtype) value = numpy.empty(shape, value.dtype)
feed.fill(c) value.fill(scalar)
else: else:
# Case 2: Reshape directly # Case 2: Reshape directly
if verify_shape: if verify_shape:
if shape is not None: if shape is not None:
if len(shape) != len(value.shape): if len(shape) != len(value.shape):
raise RuntimeError( raise RuntimeError(
'The constant was limited to {} dimensions, \ 'The constant was limited to {} dimensions, ' \
while feed a value with {} dimensions.'. 'while feed a value with {} dimensions.'
format(len(shape), len(value.shape))) .format(len(shape), len(value.shape)))
for i in range(len(shape)): for i in range(len(shape)):
if shape[i] is None: continue if shape[i] is None: continue
if shape[i] != value.shape[i]: if shape[i] != value.shape[i]:
raise RuntimeError( raise RuntimeError(
'The shape of constant was limited as (' + 'The shape of constant was limited as (' +
','.join([str(dim) for dim in shape]) + '), ' + ','.join([str(dim) for dim in shape]) + '), ' +
'while feed a value with (' + ','.join([str(dim) for dim in value.shape]) + ').') 'while feed a value with (' +
feed = feed.reshape(shape) ','.join([str(dim) for dim in value.shape]) + ').')
value = value.reshape(shape)
defined_name = dragon.workspace.GetDummyName( # Get a available name
dragon.get_default_name_scope() + defined_name = \
(name if name else 'Const'), _workspace.GetDummyName(
suffix=':0', domain='Tensor') basename=_scope.get_default_name_scope() +
(name if name else 'Const'),
suffix=':0', domain='Tensor')
# Feed into the workspace # Feed into the workspace
tensor = dragon.Tensor.Ref( return _Tensor.Ref(
name=defined_name, name=defined_name,
shape=list(feed.shape), shape=list(value.shape),
dtype=str(feed.dtype)) dtype=str(value.dtype)
tensor.set_value(feed) ).set_value(value)
return tensor \ No newline at end of file
\ No newline at end of file
Dragon/python/dragon/vm/tensorflow/framework/framework_lib.py
...@@ -13,9 +13,10 @@ from __future__ import absolute_import ...@@ -13,9 +13,10 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
from dragon.vm.tensorflow.framework.ops import Graph # The Graph (Workspace:))
from dragon.core.workspace import Workspace as Graph
# Utilities used when building a Graph. # Utilities used when building a Graph
from dragon.vm.tensorflow.framework.ops import device from dragon.vm.tensorflow.framework.ops import device
from dragon.vm.tensorflow.framework.ops import name_scope from dragon.vm.tensorflow.framework.ops import name_scope
from dragon.vm.tensorflow.framework.ops import get_default_graph from dragon.vm.tensorflow.framework.ops import get_default_graph
...@@ -27,5 +28,6 @@ from dragon.vm.tensorflow.framework.ops import GraphKeys ...@@ -27,5 +28,6 @@ from dragon.vm.tensorflow.framework.ops import GraphKeys
from dragon.vm.tensorflow.framework.constant_op import * from dragon.vm.tensorflow.framework.constant_op import *
from dragon.vm.tensorflow.framework.dtypes import * from dragon.vm.tensorflow.framework.dtypes import *
# Utilities used to represent a Tensor
from dragon.vm.tensorflow.framework.tensor_shape import Dimension from dragon.vm.tensorflow.framework.tensor_shape import Dimension
from dragon.vm.tensorflow.framework.tensor_shape import TensorShape from dragon.vm.tensorflow.framework.tensor_shape import TensorShape
\ No newline at end of file
Dragon/python/dragon/vm/tensorflow/framework/ops.py
...@@ -17,16 +17,20 @@ from __future__ import absolute_import ...@@ -17,16 +17,20 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import re from dragon.core import tls as _tls
import uuid from dragon.core import scope as _scope
import threading from dragon.core import workspace as _workspace
import dragon from dragon.core.tensor import Tensor as _Tensor
from dragon.vm.tensorflow.framework import constant_op from dragon.vm.tensorflow.framework import constant_op
from dragon.vm.tensorflow.util import tf_contextlib
def convert_to_tensor(value, dtype=None, name=None, preferred_dtype=None): def convert_to_tensor(
value,
dtype=None,
name=None,
preferred_dtype=None,
):
"""Converts the given value to a Tensor. """Converts the given value to a Tensor.
Parameters Parameters
...@@ -46,73 +50,10 @@ def convert_to_tensor(value, dtype=None, name=None, preferred_dtype=None): ...@@ -46,73 +50,10 @@ def convert_to_tensor(value, dtype=None, name=None, preferred_dtype=None):
The output tensor. The output tensor.
""" """
if isinstance(value, dragon.Tensor): return value if isinstance(value, _Tensor): return value
return constant_op.constant(value, dtype=dtype, name=name) return constant_op.constant(value, dtype=dtype, name=name)
class Graph(object):
"""A wrapper to connect ``Function`` to ``Workspace``.
Note that official TensorFlow trace the expressions explicitly
in this class, while we have done in the virtual stack.
Besides, organizing a ``Flow``, i.e., expressions with specified
outputs should also be done here.
"""
def __init__(self):
self._collections = {}
self._workspace = 'tf/graph/' + str(uuid.uuid4())
def get_collection_ref(self, name):
coll_list = self._collections.get(name, None)
if coll_list is None:
coll_list = []
self._collections[name] = coll_list
return coll_list
def get_collection(self, name, scope=None):
coll_list = self._collections.get(name, None)
if coll_list is None:
return []
if scope is None:
return list(coll_list)
else:
filter_coll_list = []
regex = re.compile(scope)
for item in coll_list:
if hasattr(item, "name") and regex.match(item.name):
filter_coll_list.append(item)
return filter_coll_list
def add_to_collection(self, name, value):
if name not in self._collections:
self._collections[name] = [value]
else:
self._collections[name].append(value)
def add_to_collections(self, names, value):
for name in names:
self.add_to_collection(name, value)
def device(self, device_name_or_function):
if not isinstance(device_name_or_function, str):
raise TypeError('The device function should be a str.')
device_and_id = device_name_or_function.split('/')[1]
device, id = device_and_id.split(':')
if device not in ['cpu', 'gpu']:
raise ValueError('The device should either be cpu or gpu.')
try:
id = int(id)
except Exception as e:
raise ValueError('The device id should be a integer.')
return dragon.device_scope(device, device_id=id)
def as_default(self):
return _default_graph_stack.get_controller(self)
class GraphKeys(object): class GraphKeys(object):
GLOBAL_VARIABLES = "variables" GLOBAL_VARIABLES = "variables"
# Key to collect local variables that are local to the machine and are not # Key to collect local variables that are local to the machine and are not
...@@ -202,112 +143,15 @@ def add_to_collections(names, value): ...@@ -202,112 +143,15 @@ def add_to_collections(names, value):
def name_scope(name, default_name=None, values=None): def name_scope(name, default_name=None, values=None):
name = default_name if name is None else name name = default_name if name is None else name
name = '' if name is None else name name = '' if name is None else name
return dragon.name_scope(name) return _scope.name_scope(name)
##############################################
# #
# Default Stack #
# #
##############################################
class _DefaultStack(threading.local):
"""A thread-local stack of objects for providing implicit defaults."""
def __init__(self):
super(_DefaultStack, self).__init__()
self._enforce_nesting = True
self.stack = []
def get_default(self):
return self.stack[-1] if len(self.stack) >= 1 else None
def reset(self):
self.stack = []
def is_cleared(self):
return not self.stack
@property
def enforce_nesting(self):
return self._enforce_nesting
@enforce_nesting.setter
def enforce_nesting(self, value):
self._enforce_nesting = value
@tf_contextlib.contextmanager
def get_controller(self, default):
"""A context manager for manipulating a default stack."""
self.stack.append(default)
try:
yield default
finally:
# stack may be empty if reset() was called
if self.stack:
if self._enforce_nesting:
if self.stack[-1] is not default:
raise AssertionError(
"Nesting violated for default stack of %s objects" %
type(default))
self.stack.pop()
else:
self.stack.remove(default)
class _DefaultGraphStack(_DefaultStack):
"""A thread-local stack of objects for providing an implicit default graph."""
def __init__(self):
super(_DefaultGraphStack, self).__init__()
self._global_default_graph = None
def get_default(self):
"""Override that returns a global default if the stack is empty."""
ret = super(_DefaultGraphStack, self).get_default()
if ret is None:
ret = self._GetGlobalDefaultGraph()
return ret
def _GetGlobalDefaultGraph(self):
if self._global_default_graph is None:
# TODO(mrry): Perhaps log that the default graph is being used, or set
# provide some other feedback to prevent confusion when a mixture of
# the global default graph and an explicit graph are combined in the
# same process.
self._global_default_graph = Graph()
# Rewritten the random workspace name
self._global_default_graph._workspace = 'default'
return self._global_default_graph
def reset(self):
super(_DefaultGraphStack, self).reset()
# We should call dragon api to reset the workspace
dragon.workspace.ResetWorkspace(self._global_default_graph._workspace)
self._global_default_graph = None
@tf_contextlib.contextmanager
def get_controller(self, default):
with super(_DefaultGraphStack, self).get_controller(default) as g:
with dragon.ws_scope(g._workspace):
yield g
_default_graph_stack = _DefaultGraphStack()
_default_session_stack = _DefaultStack()
def get_default_graph(): def get_default_graph():
return _default_graph_stack.get_default() return _workspace.get_default_workspace()
def reset_default_graph(): def reset_default_graph():
if not _default_graph_stack.is_cleared(): _workspace.reset_default_workspace()
raise AssertionError("Do not use tf.reset_default_graph() to clear "
"nested graphs. If you need a cleared graph, "
"exit the nesting and create a new graph.")
_default_graph_stack.reset()
def default_session(session): def default_session(session):
...@@ -319,7 +163,17 @@ def get_default_session(): ...@@ -319,7 +163,17 @@ def get_default_session():
def device(device_name_or_function): def device(device_name_or_function):
return get_default_graph().device(device_name_or_function) if not isinstance(device_name_or_function, str):
raise TypeError('The device function should be a str.')
device_and_id = device_name_or_function.split('/')[1]
device, id = device_and_id.split(':')
if device not in ['cpu', 'gpu']:
raise ValueError('The device should either be cpu or gpu.')
try:
id = int(id)
except Exception as _:
raise ValueError('The device id should be a integer.')
return _scope.device_scope(device, device_id=id)
def _eval_using_default_session(tensors, feed_dict, session=None): def _eval_using_default_session(tensors, feed_dict, session=None):
...@@ -333,6 +187,10 @@ def _eval_using_default_session(tensors, feed_dict, session=None): ...@@ -333,6 +187,10 @@ def _eval_using_default_session(tensors, feed_dict, session=None):
return session.run(tensors, feed_dict) return session.run(tensors, feed_dict)
_default_session_stack = _tls.Stack()
# The Monkey Patching
# Require "import dragon.vm.tensorflow" # Require "import dragon.vm.tensorflow"
dragon.Tensor.eval = lambda self, feed_dict=None, session=None : \ _Tensor.eval = lambda self, feed_dict=None, session=None : \
_eval_using_default_session(self, feed_dict, session) _eval_using_default_session(self, feed_dict, session)
\ No newline at end of file
Dragon/python/dragon/vm/tensorflow/framework/tensor_shape.py
...@@ -13,7 +13,7 @@ from __future__ import absolute_import ...@@ -13,7 +13,7 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
from dragon.core.tensor import Tensor from dragon.core.tensor import Tensor as _Tensor
class Dimension(object): class Dimension(object):
...@@ -114,4 +114,5 @@ def get_shape(self): ...@@ -114,4 +114,5 @@ def get_shape(self):
return TensorShape(self.shape) return TensorShape(self.shape)
Tensor.get_shape = get_shape # The Monkey Patching
\ No newline at end of file _Tensor.get_shape = get_shape
\ No newline at end of file
Dragon/python/dragon/vm/tensorflow/layers/base.py
...@@ -28,7 +28,13 @@ from dragon.vm.tensorflow.util import nest ...@@ -28,7 +28,13 @@ from dragon.vm.tensorflow.util import nest
class Layer(object): class Layer(object):
def __init__(self, trainable=True, name=None, dtype=dtypes.float32, **kwargs): def __init__(
self,
trainable=True,
name=None,
dtype=dtypes.float32,
**kwargs
):
allowed_kwargs = {'_scope', '_reuse'} allowed_kwargs = {'_scope', '_reuse'}
for kwarg in kwargs: for kwarg in kwargs:
if kwarg not in allowed_kwargs: if kwarg not in allowed_kwargs:
...@@ -79,13 +85,15 @@ class Layer(object): ...@@ -79,13 +85,15 @@ class Layer(object):
_add_elements_to_collection(self.updates, ops.GraphKeys.UPDATE_OPS) _add_elements_to_collection(self.updates, ops.GraphKeys.UPDATE_OPS)
return outputs return outputs
def add_variable(self, def add_variable(
name, self,
shape, name,
dtype=None, shape,
trainable=True, dtype=None,
initializer=None, trainable=True,
regularizer=None): initializer=None,
regularizer=None,
):
if dtype is None: dtype = self.dtype if dtype is None: dtype = self.dtype
variable = vs.get_variable( variable = vs.get_variable(
name, name,
...@@ -93,7 +101,8 @@ class Layer(object): ...@@ -93,7 +101,8 @@ class Layer(object):
initializer=initializer, initializer=initializer,
regularizer=regularizer, regularizer=regularizer,
dtype=dtypes.as_dtype(dtype), dtype=dtypes.as_dtype(dtype),
trainable=trainable and self.trainable) trainable=trainable and self.trainable,
)
if trainable: if trainable:
self._trainable_weights.append(variable) self._trainable_weights.append(variable)
else: else:
...@@ -105,9 +114,14 @@ class Layer(object): ...@@ -105,9 +114,14 @@ class Layer(object):
class InputSpec(object): class InputSpec(object):
def __init__(self, def __init__(
dtype=None, shape=None, ndim=None, self,
max_ndim=None, min_ndim=None, axes=None dtype=None,
shape=None,
ndim=None,
max_ndim=None,
min_ndim=None,
axes=None,
): ):
self.dtype = dtype self.dtype = dtype
self.shape = shape self.shape = shape
...@@ -125,9 +139,6 @@ def _to_snake_case(name): ...@@ -125,9 +139,6 @@ def _to_snake_case(name):
return 'private' + insecure return 'private' + insecure
PER_GRAPH_LAYER_NAME_UIDS = weakref.WeakKeyDictionary()
def _unique_layer_name(name): def _unique_layer_name(name):
global PER_GRAPH_LAYER_NAME_UIDS global PER_GRAPH_LAYER_NAME_UIDS
graph = ops.get_default_graph() graph = ops.get_default_graph()
...@@ -153,3 +164,6 @@ def _add_elements_to_collection(elements, collection_list): ...@@ -153,3 +164,6 @@ def _add_elements_to_collection(elements, collection_list):
for element in elements: for element in elements:
if element not in collection_set: if element not in collection_set:
collection.append(element) collection.append(element)
PER_GRAPH_LAYER_NAME_UIDS = weakref.WeakKeyDictionary()
\ No newline at end of file
Dragon/python/dragon/vm/tensorflow/layers/convolutional.py
...@@ -20,24 +20,26 @@ from dragon.vm.tensorflow.ops import nn ...@@ -20,24 +20,26 @@ from dragon.vm.tensorflow.ops import nn
class _Conv(base.Layer): class _Conv(base.Layer):
def __init__(self, def __init__(
rank, self,
filters, rank,
kernel_size, filters,
strides=1, kernel_size,
padding='valid', strides=1,
data_format='channels_last', padding='valid',
dilation_rate=1, data_format='channels_last',
activation=None, dilation_rate=1,
use_bias=True, activation=None,
kernel_initializer=None, use_bias=True,
bias_initializer=init_ops.zeros_initializer(), kernel_initializer=None,
kernel_regularizer=None, bias_initializer=init_ops.zeros_initializer(),
bias_regularizer=None, kernel_regularizer=None,
activity_regularizer=None, bias_regularizer=None,
trainable=True, activity_regularizer=None,
name=None, trainable=True,
**kwargs): name=None,
**kwargs
):
super(_Conv, self).__init__(trainable=trainable, name=name, **kwargs) super(_Conv, self).__init__(trainable=trainable, name=name, **kwargs)
self.rank = rank self.rank = rank
self.filters = filters self.filters = filters
...@@ -82,7 +84,8 @@ class _Conv(base.Layer): ...@@ -82,7 +84,8 @@ class _Conv(base.Layer):
shape=kernel_shape, shape=kernel_shape,
initializer=self.kernel_initializer, initializer=self.kernel_initializer,
regularizer=self.kernel_regularizer, regularizer=self.kernel_regularizer,
dtype=self.dtype) dtype=self.dtype,
)
if self.use_bias: if self.use_bias:
self.bias = self.add_variable( self.bias = self.add_variable(
...@@ -90,7 +93,8 @@ class _Conv(base.Layer): ...@@ -90,7 +93,8 @@ class _Conv(base.Layer):
shape=(self.filters,), shape=(self.filters,),
initializer=self.bias_initializer, initializer=self.bias_initializer,
regularizer=self.bias_regularizer, regularizer=self.bias_regularizer,
dtype=self.dtype) dtype=self.dtype,
)
else: else:
self.bias = None self.bias = None
...@@ -108,10 +112,15 @@ class _Conv(base.Layer): ...@@ -108,10 +112,15 @@ class _Conv(base.Layer):
dilation_rate=self.dilation_rate, dilation_rate=self.dilation_rate,
strides=self.strides, strides=self.strides,
padding=self.padding.upper(), padding=self.padding.upper(),
data_format=tf_data_format) data_format=tf_data_format,
)
if self.bias is not None: if self.bias is not None:
outputs = nn.bias_add(outputs, self.bias, data_format=tf_data_format) outputs = nn.bias_add(
outputs,
self.bias,
data_format=tf_data_format,
)
if self.activation is not None: if self.activation is not None:
return self.activation(outputs) return self.activation(outputs)
...@@ -119,22 +128,25 @@ class _Conv(base.Layer): ...@@ -119,22 +128,25 @@ class _Conv(base.Layer):
class Conv2D(_Conv): class Conv2D(_Conv):
def __init__(self, filters, def __init__(
kernel_size, self,
strides=(1, 1), filters,
padding='valid', kernel_size,
data_format='channels_last', strides=(1, 1),
dilation_rate=(1, 1), padding='valid',
activation=None, data_format='channels_last',
use_bias=True, dilation_rate=(1, 1),
kernel_initializer=None, activation=None,
bias_initializer=init_ops.zeros_initializer(), use_bias=True,
kernel_regularizer=None, kernel_initializer=None,
bias_regularizer=None, bias_initializer=init_ops.zeros_initializer(),
activity_regularizer=None, kernel_regularizer=None,
trainable=True, bias_regularizer=None,
name=None, activity_regularizer=None,
**kwargs): trainable=True,
name=None,
**kwargs
):
super(Conv2D, self).__init__( super(Conv2D, self).__init__(
rank=2, rank=2,
filters=filters, filters=filters,
...@@ -154,23 +166,25 @@ class Conv2D(_Conv): ...@@ -154,23 +166,25 @@ class Conv2D(_Conv):
name=name, **kwargs) name=name, **kwargs)
def conv2d(inputs, def conv2d(
filters, inputs,
kernel_size, filters,
strides=(1, 1), kernel_size,
padding='valid', strides=(1, 1),
data_format='channels_last', padding='valid',
dilation_rate=(1, 1), data_format='channels_last',
activation=None, dilation_rate=(1, 1),
use_bias=True, activation=None,
kernel_initializer=None, use_bias=True,
bias_initializer=init_ops.zeros_initializer(), kernel_initializer=None,
kernel_regularizer=None, bias_initializer=init_ops.zeros_initializer(),
bias_regularizer=None, kernel_regularizer=None,
activity_regularizer=None, bias_regularizer=None,
trainable=True, activity_regularizer=None,
name=None, trainable=True,
reuse=None): name=None,
reuse=None,
):
return Conv2D( return Conv2D(
filters=filters, filters=filters,
kernel_size=kernel_size, kernel_size=kernel_size,
...@@ -188,4 +202,5 @@ def conv2d(inputs, ...@@ -188,4 +202,5 @@ def conv2d(inputs,
trainable=trainable, trainable=trainable,
name=name, name=name,
_reuse=reuse, _reuse=reuse,
_scope=name).apply(inputs) _scope=name,
\ No newline at end of file ).apply(inputs)
\ No newline at end of file
Dragon/python/dragon/vm/tensorflow/layers/core.py
...@@ -24,18 +24,20 @@ from dragon.vm.tensorflow.ops import standard_ops ...@@ -24,18 +24,20 @@ from dragon.vm.tensorflow.ops import standard_ops
class Dense(base.Layer): class Dense(base.Layer):
def __init__(self, def __init__(
units, self,
activation=None, units,
use_bias=True, activation=None,
kernel_initializer=None, use_bias=True,
bias_initializer=init_ops.zeros_initializer(), kernel_initializer=None,
kernel_regularizer=None, bias_initializer=init_ops.zeros_initializer(),
bias_regularizer=None, kernel_regularizer=None,
activity_regularizer=None, bias_regularizer=None,
trainable=True, activity_regularizer=None,
name=None, trainable=True,
**kwargs): name=None,
**kwargs
):
super(Dense, self).__init__(trainable=trainable, name=name, **kwargs) super(Dense, self).__init__(trainable=trainable, name=name, **kwargs)
self.units = units self.units = units
self.activation = activation self.activation = activation
...@@ -61,7 +63,8 @@ class Dense(base.Layer): ...@@ -61,7 +63,8 @@ class Dense(base.Layer):
shape=[input_shape[-1].value, self.units], shape=[input_shape[-1].value, self.units],
initializer=self.kernel_initializer, initializer=self.kernel_initializer,
regularizer=self.kernel_regularizer, regularizer=self.kernel_regularizer,
dtype=self.dtype) dtype=self.dtype,
)
if self.use_bias: if self.use_bias:
self.bias = self.add_variable( self.bias = self.add_variable(
...@@ -69,7 +72,8 @@ class Dense(base.Layer): ...@@ -69,7 +72,8 @@ class Dense(base.Layer):
shape=[self.units,], shape=[self.units,],
initializer=self.bias_initializer, initializer=self.bias_initializer,
regularizer=self.bias_regularizer, regularizer=self.bias_regularizer,
dtype=self.dtype) dtype=self.dtype,
)
else: else:
self.bias = None self.bias = None
self.built = True self.built = True
...@@ -83,18 +87,20 @@ class Dense(base.Layer): ...@@ -83,18 +87,20 @@ class Dense(base.Layer):
return outputs return outputs
def dense(inputs, def dense(
units, inputs,
activation=None, units,
use_bias=True, activation=None,
kernel_initializer=None, use_bias=True,
bias_initializer=init_ops.zeros_initializer(), kernel_initializer=None,
kernel_regularizer=None, bias_initializer=init_ops.zeros_initializer(),
bias_regularizer=None, kernel_regularizer=None,
activity_regularizer=None, bias_regularizer=None,
trainable=True, activity_regularizer=None,
name=None, trainable=True,
reuse=None): name=None,
reuse=None,
):
return Dense( return Dense(
units, units,
activation=activation, activation=activation,
...@@ -107,4 +113,5 @@ def dense(inputs, ...@@ -107,4 +113,5 @@ def dense(inputs,
trainable=trainable, trainable=trainable,
name=name, name=name,
_scope=name, _scope=name,
_reuse=reuse).apply(inputs) _reuse=reuse,
\ No newline at end of file ).apply(inputs)
\ No newline at end of file
Dragon/python/dragon/vm/tensorflow/layers/layers.py
...@@ -13,20 +13,20 @@ from __future__ import absolute_import ...@@ -13,20 +13,20 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
from dragon.vm.tensorflow.layers.convolutional import ( from .convolutional import (
conv2d, Conv2D, conv2d, Conv2D,
) )
from dragon.vm.tensorflow.layers.core import ( from .core import (
dense, Dense, dense, Dense,
) )
from dragon.vm.tensorflow.layers.normalization import ( from .normalization import (
batch_normalization, BatchNormalization, batch_normalization, BatchNormalization,
batch_norm, BatchNorm, batch_norm, BatchNorm,
) )
from dragon.vm.tensorflow.layers.pooling import ( from .pooling import (
average_pooling2d, AveragePooling2D, average_pooling2d, AveragePooling2D,
max_pooling2d, MaxPooling2D, max_pooling2d, MaxPooling2D,
) )
Dragon/python/dragon/vm/tensorflow/layers/normalization.py
...@@ -17,7 +17,7 @@ from __future__ import absolute_import ...@@ -17,7 +17,7 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import dragon from dragon.ops import BatchNorm as _BatchNormOp
from dragon.vm.tensorflow.framework import tensor_shape from dragon.vm.tensorflow.framework import tensor_shape
from dragon.vm.tensorflow.layers import base from dragon.vm.tensorflow.layers import base
...@@ -25,26 +25,29 @@ from dragon.vm.tensorflow.ops import init_ops ...@@ -25,26 +25,29 @@ from dragon.vm.tensorflow.ops import init_ops
class BatchNormalization(base.Layer): class BatchNormalization(base.Layer):
def __init__(self, def __init__(
axis=-1, self,
momentum=0.99, axis=-1,
epsilon=1e-3, momentum=0.99,
center=True, epsilon=1e-3,
scale=True, center=True,
beta_initializer=init_ops.zeros_initializer(), scale=True,
gamma_initializer=init_ops.ones_initializer(), beta_initializer=init_ops.zeros_initializer(),
moving_mean_initializer=init_ops.zeros_initializer(), gamma_initializer=init_ops.ones_initializer(),
moving_variance_initializer=init_ops.ones_initializer(), moving_mean_initializer=init_ops.zeros_initializer(),
beta_regularizer=None, moving_variance_initializer=init_ops.ones_initializer(),
gamma_regularizer=None, beta_regularizer=None,
renorm=False, gamma_regularizer=None,
renorm_clipping=None, renorm=False,
renorm_momentum=0.99, renorm_clipping=None,
fused=None, renorm_momentum=0.99,
trainable=True, fused=None,
name=None, trainable=True,
**kwargs): name=None,
super(BatchNormalization, self).__init__(trainable=trainable, name=name, **kwargs) **kwargs
):
super(BatchNormalization, self).__init__(
trainable=trainable, name=name, **kwargs)
self.axis = axis self.axis = axis
self.momentum = momentum self.momentum = momentum
self.epsilon = epsilon self.epsilon = epsilon
...@@ -92,33 +95,37 @@ class BatchNormalization(base.Layer): ...@@ -92,33 +95,37 @@ class BatchNormalization(base.Layer):
name='moving_mean', name='moving_mean',
shape=(param_dim.value,), shape=(param_dim.value,),
initializer=self.moving_mean_initializer, initializer=self.moving_mean_initializer,
trainable=False) trainable=False,
)
self.moving_variance = self.add_variable( self.moving_variance = self.add_variable(
name='moving_variance', name='moving_variance',
shape=(param_dim.value,), shape=(param_dim.value,),
initializer=self.moving_variance_initializer, initializer=self.moving_variance_initializer,
trainable=False) trainable=False,
)
self.gamma = self.add_variable( self.gamma = self.add_variable(
name='gamma', name='gamma',
shape=(param_dim.value,), shape=(param_dim.value,),
initializer=self.gamma_initializer, initializer=self.gamma_initializer,
regularizer=self.gamma_regularizer, regularizer=self.gamma_regularizer,
trainable=self.scale) trainable=self.scale,
)
self.beta = self.add_variable( self.beta = self.add_variable(
name='beta', name='beta',
shape=(param_dim.value,), shape=(param_dim.value,),
initializer=self.beta_initializer, initializer=self.beta_initializer,
regularizer=self.beta_regularizer, regularizer=self.beta_regularizer,
trainable=self.center) trainable=self.center,
)
self.built = True self.built = True
def call(self, inputs, training=False, *args, **kwargs): def call(self, inputs, training=False, *args, **kwargs):
use_stats = 0 if training else 1 use_stats = 0 if training else 1
return dragon.ops.BatchNorm([ return _BatchNormOp([
inputs, inputs,
self.moving_mean, self.moving_mean,
self.moving_variance, self.moving_variance,
...@@ -127,7 +134,8 @@ class BatchNormalization(base.Layer): ...@@ -127,7 +134,8 @@ class BatchNormalization(base.Layer):
axis=self.axis, axis=self.axis,
momentum=self.momentum, momentum=self.momentum,
eps=self.epsilon, eps=self.epsilon,
use_stats=use_stats) use_stats=use_stats,
)
def batch_normalization( def batch_normalization(
...@@ -170,7 +178,8 @@ def batch_normalization( ...@@ -170,7 +178,8 @@ def batch_normalization(
trainable=trainable, trainable=trainable,
name=name, name=name,
_reuse=reuse, _reuse=reuse,
_scope=name).apply(inputs, training=training) _scope=name,
).apply(inputs, training=training)
# Aliases # Aliases
......
Dragon/python/dragon/vm/tensorflow/layers/pooling.py
...@@ -22,9 +22,16 @@ from dragon.vm.tensorflow.layers import base, utils ...@@ -22,9 +22,16 @@ from dragon.vm.tensorflow.layers import base, utils
class _Pooling2D(base.Layer): class _Pooling2D(base.Layer):
def __init__(self, pool_function, pool_size, strides, def __init__(
padding='valid', data_format='channels_last', self,
name=None, **kwargs): pool_function,
pool_size,
strides,
padding='valid',
data_format='channels_last',
name=None,
**kwargs
):
super(_Pooling2D, self).__init__(name=name, **kwargs) super(_Pooling2D, self).__init__(name=name, **kwargs)
self.pool_function = pool_function self.pool_function = pool_function
self.pool_size = utils.normalize_tuple(pool_size, 2, 'pool_size') self.pool_size = utils.normalize_tuple(pool_size, 2, 'pool_size')
...@@ -40,19 +47,25 @@ class _Pooling2D(base.Layer): ...@@ -40,19 +47,25 @@ class _Pooling2D(base.Layer):
else: else:
pool_shape = (1, 1) + self.pool_size pool_shape = (1, 1) + self.pool_size
strides = (1, 1) + self.strides strides = (1, 1) + self.strides
outputs = self.pool_function( return self.pool_function(
inputs, inputs,
ksize=pool_shape, ksize=pool_shape,
strides=strides, strides=strides,
padding=self.padding.upper(), padding=self.padding.upper(),
data_format=utils.convert_data_format(self.data_format, 4)) data_format=utils.convert_data_format(self.data_format, 4),
return outputs )
class MaxPooling2D(_Pooling2D): class MaxPooling2D(_Pooling2D):
def __init__(self, pool_size, strides, def __init__(
padding='valid', data_format='channels_last', self,
name=None, **kwargs): pool_size,
strides,
padding='valid',
data_format='channels_last',
name=None,
**kwargs
):
super(MaxPooling2D, self).__init__( super(MaxPooling2D, self).__init__(
nn.max_pool, nn.max_pool,
pool_size=pool_size, pool_size=pool_size,
...@@ -63,9 +76,15 @@ class MaxPooling2D(_Pooling2D): ...@@ -63,9 +76,15 @@ class MaxPooling2D(_Pooling2D):
class AveragePooling2D(_Pooling2D): class AveragePooling2D(_Pooling2D):
def __init__(self, pool_size, strides, def __init__(
padding='valid', data_format='channels_last', self,
name=None, **kwargs): pool_size,
strides,
padding='valid',
data_format='channels_last',
name=None,
**kwargs
):
super(AveragePooling2D, self).__init__( super(AveragePooling2D, self).__init__(
nn.avg_pool, nn.avg_pool,
pool_size=pool_size, pool_size=pool_size,
...@@ -76,22 +95,34 @@ class AveragePooling2D(_Pooling2D): ...@@ -76,22 +95,34 @@ class AveragePooling2D(_Pooling2D):
def max_pooling2d( def max_pooling2d(
inputs, pool_size, strides, padding='valid', inputs,
data_format='channels_last', name=None): pool_size,
strides,
padding='valid',
data_format='channels_last',
name=None,
):
return MaxPooling2D( return MaxPooling2D(
pool_size=pool_size, pool_size=pool_size,
strides=strides, strides=strides,
padding=padding, padding=padding,
data_format=data_format, data_format=data_format,
name=name).apply(inputs) name=name,
).apply(inputs)
def average_pooling2d( def average_pooling2d(
inputs, pool_size, strides, padding='valid', inputs,
data_format='channels_last', name=None): pool_size,
strides,
padding='valid',
data_format='channels_last',
name=None,
):
return AveragePooling2D( return AveragePooling2D(
pool_size=pool_size, pool_size=pool_size,
strides=strides, strides=strides,
padding=padding, padding=padding,
data_format=data_format, data_format=data_format,
name=name).apply(inputs) name=name,
\ No newline at end of file ).apply(inputs)
\ No newline at end of file
Markdown is supported
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!