Skip to content

SeetaResearch / Dragon

Go to a project
Commit bab9931e by Ting PAN
Options

PyTorch + TorchVision + Dragon = Dragon

1 parent 176b7bbb
Expand all
Inline Side-by-side
Showing with 3407 additions and 398 deletions
Docker/ubuntu-16.04-cpu-openblas/Dockerfile
......@@ -8,6 +8,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
unzip \
ssh \
vim \
libnuma-dev \
libprotobuf-dev \
protobuf-compiler \
libopenblas-dev \
......@@ -17,7 +18,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
python3-tk \
&& rm -rf /var/lib/apt/lists/*
RUN pip3 install --no-cache-dir --upgrade pip setuptools wheel && \
RUN pip3 install --no-cache-dir --upgrade setuptools wheel && \
pip3 install --no-cache-dir -i https://pypi.tuna.tsinghua.edu.cn/simple \
numpy \
protobuf \
......
Docker/ubuntu-16.04-cuda9.0-cudnn7/Dockerfile
......@@ -8,6 +8,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
unzip \
ssh \
vim \
libnuma-dev \
libprotobuf-dev \
protobuf-compiler \
libopenblas-dev \
......@@ -19,7 +20,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
python3-tk \
&& rm -rf /var/lib/apt/lists/*
RUN pip3 install --no-cache-dir --upgrade pip setuptools wheel && \
RUN pip3 install --no-cache-dir --upgrade setuptools wheel && \
pip3 install --no-cache-dir -i https://pypi.tuna.tsinghua.edu.cn/simple \
numpy \
protobuf \
......
Dragon/include/core/operator.h
......@@ -68,6 +68,7 @@ class OperatorBase {
inline const OperatorDef& op_def() const { return op_def_; }
inline string DebugString() const { return op_def_.DebugString(); }
string DTypeHelper(const Tensor& tensor, const Set<string>& dtypes) const;
string DTypeHelper(const string& dtype, const Set<string>& dtypes) const;
protected:
string phase_, anchor_;
......
Dragon/include/core/tensor.h
......@@ -206,9 +206,8 @@ class Tensor {
template <class DstCTX, class SrcCTX>
inline void Copy(const Tensor& other) {
CHECK_EQ(size_, other.size_);
meta_ = other.meta_;
auto* src = other.template raw_data<SrcCTX>();
auto* dst = raw_mutable_data<DstCTX>();
auto* dst = raw_mutable_data<DstCTX>(other.meta_);
if (dst == src) return;
if (TypeMeta::Id<DstCTX>() == TypeMeta::Id<CPUContext>()) {
CPUContext::Memcpy<DstCTX, SrcCTX>(nbytes(), dst, src);
......
Dragon/include/operators/activation/dropout_op.h
......@@ -33,7 +33,6 @@ class DropoutOp final : public Operator<Context> {
protected:
DECLARE_ARGUMENT_WITH_DESC(float, prob);
bool use_scale;
Tensor* mask;
};
template <class Context>
......
Dragon/include/operators/cast/float2half_op.h → Dragon/include/operators/misc/astype_op.h
......@@ -9,21 +9,29 @@
//
// ------------------------------------------------------------
#ifndef DRAGON_OPERATORS_CAST_FLOAT2HALF_OP_H_
#define DRAGON_OPERATORS_CAST_FLOAT2HALF_OP_H_
#ifndef DRAGON_OPERATORS_MISC_ASTYPE_OP_H_
#define DRAGON_OPERATORS_MISC_ASTYPE_OP_H_
#include "core/operator.h"
namespace dragon {
template <class Context>
class FloatToHalfOp final : public Operator<Context> {
class AsTypeOp final : public Operator<Context> {
public:
USE_SIMPLE_CTOR_DTOR(FloatToHalfOp);
AsTypeOp(const OperatorDef& op_def, Workspace* ws)
: Operator<Context>(op_def, ws),
dtype(OperatorBase::GetSingleArg<string>("dtype", "float32")),
inplace(OperatorBase::GetSingleArg<bool>("inplace", false)) {}
USE_OPERATOR_FUNCTIONS(Context);
void RunOnDevice() override;
protected:
string dtype;
bool inplace;
};
} // namespace dragon
#endif // DRAGON_OPERATORS_CAST_FLOAT2HALF_OP_H_
\ No newline at end of file
#endif // DRAGON_OPERATORS_MISC_ASTYPE_OP_H_
\ No newline at end of file
Dragon/include/operators/ndarray/argmin_op.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_OPERATORS_NDARRAY_ARGMIN_OP_H_
#define DRAGON_OPERATORS_NDARRAY_ARGMIN_OP_H_
#include "core/operator.h"
namespace dragon {
template <class Context>
class ArgminOp final : public Operator<Context> {
public:
ArgminOp(const OperatorDef& op_def, Workspace* ws)
: Operator<Context>(op_def, ws),
axis(OperatorBase::GetSingleArg<int>("axis", -1)),
keep_dims(OperatorBase::GetSingleArg<bool>("keep_dims", false)),
top_k(OperatorBase::GetSingleArg<int>("top_k", 1)) {}
USE_OPERATOR_FUNCTIONS(Context);
void RunOnDevice() override;
template <typename T> void RunWithType();
protected:
TIndex axis, axis_dim, top_k, count, inner_dim;
bool keep_dims;
};
} // namespace dragon
#endif // DRAGON_OPERATORS_NDARRAY_ARGMIN_OP_H_
\ No newline at end of file
Dragon/include/operators/ndarray/argmax_op.h → Dragon/include/operators/ndarray/argreduce_op.h
......@@ -9,19 +9,20 @@
//
// -------------------------------------------------------------
#ifndef DRAGON_OPERATORS_NDARRAY_ARGMAX_OP_H_
#define DRAGON_OPERATORS_NDARRAY_ARGMAX_OP_H_
#ifndef DRAGON_OPERATORS_NDARRAY_ARGREDUCE_OP_H_
#define DRAGON_OPERATORS_NDARRAY_ARGREDUCE_OP_H_
#include "core/operator.h"
namespace dragon {
template <class Context>
class ArgmaxOp final : public Operator<Context> {
class ArgReduceOp final : public Operator<Context> {
public:
ArgmaxOp(const OperatorDef& op_def, Workspace* ws)
ArgReduceOp(const OperatorDef& op_def, Workspace* ws)
: Operator<Context>(op_def, ws),
axis(OperatorBase::GetSingleArg<int>("axis", -1)),
operation(OperatorBase::GetSingleArg<string>("operation", "NONE")),
keep_dims(OperatorBase::GetSingleArg<bool>("keep_dims", false)),
top_k(OperatorBase::GetSingleArg<int>("top_k", 1)) {}
USE_OPERATOR_FUNCTIONS(Context);
......@@ -30,10 +31,11 @@ class ArgmaxOp final : public Operator<Context> {
template <typename T> void RunWithType();
protected:
TIndex axis, axis_dim, top_k, count, inner_dim;
bool keep_dims;
string operation;
TIndex axis, axis_dim, top_k, count, inner_dim;
};
} // namespace dragon
#endif // DRAGON_OPERATORS_NDARRAY_ARGMAX_OP_H_
\ No newline at end of file
#endif // DRAGON_OPERATORS_NDARRAY_ARGREDUCE_OP_H_
\ No newline at end of file
Dragon/include/operators/ndarray/concat_op.h
......@@ -20,16 +20,15 @@ template <class Context>
class ConcatOp : public Operator<Context> {
public:
ConcatOp(const OperatorDef& op_def, Workspace* ws)
: Operator<Context>(op_def, ws),
axis(OperatorBase::GetSingleArg<int>("axis", 1)),
nin(OperatorBase::GetSingleArg<int>("num_input", 1)) {}
: Operator<Context>(op_def, ws),
axis(OperatorBase::GetSingleArg<int>("axis", 1)) {}
USE_OPERATOR_FUNCTIONS(Context);
void RunOnDevice() override;
template <typename T> void RunWithType();
protected:
TIndex axis, nin, outer_dim, inner_dim, x_concat_dim, y_concat_dim;
TIndex axis, outer_dim, inner_dim, x_concat_dim, y_concat_dim;
TIndex x_offset, y_offset, concat_offset;
vector<TIndex> concat_dims;
};
......@@ -39,15 +38,14 @@ class ConcatGradientOp : public Operator<Context> {
public:
ConcatGradientOp(const OperatorDef& op_def, Workspace* ws)
: Operator<Context>(op_def, ws),
axis(OperatorBase::GetSingleArg<int>("axis", 1)),
nin(OperatorBase::GetSingleArg<int>("num_input", 1)) {}
axis(OperatorBase::GetSingleArg<int>("axis", 1)) {}
USE_OPERATOR_FUNCTIONS(Context);
void RunOnDevice() override;
template <typename T> void RunWithType();
protected:
TIndex axis, nin, outer_dim, inner_dim, x_concat_dim, y_concat_dim;
TIndex axis, outer_dim, inner_dim, x_concat_dim, y_concat_dim;
TIndex x_offset, y_offset, concat_offset;
vector<TIndex> concat_dims;
};
......
Dragon/include/operators/norm/batch_norm_op.h
......@@ -12,6 +12,8 @@
#ifndef DRAGON_OPERATORS_NORM_BATCH_NORM_OP_H_
#define DRAGON_OPERATORS_NORM_BATCH_NORM_OP_H_
#include <cfloat>
#include "core/operator.h"
namespace dragon {
......@@ -22,8 +24,8 @@ class BatchNormOp : public Operator<Context> {
BatchNormOp(const OperatorDef& op_def, Workspace* ws)
: Operator<Context>(op_def, ws),
axis(OperatorBase::GetSingleArg<int>("axis", -1)),
momentum(OperatorBase::GetSingleArg<float>("momentum", float(0.9))),
eps(OperatorBase::GetSingleArg<float>("eps", float(1e-3))),
momentum(OperatorBase::GetSingleArg<float>("momentum", 0.9f)),
eps(OperatorBase::GetSingleArg<float>("eps", 1e-3f)),
use_stats(OperatorBase::GetSingleArg<int>("use_stats", -1)),
mode(OperatorBase::GetSingleArg<string>("mode", "DEFAULT")) {
if (axis != -1)
......@@ -84,8 +86,8 @@ class FusedBatchNormOp : public Operator<Context> {
FusedBatchNormOp(const OperatorDef& op_def, Workspace* ws)
: Operator<Context>(op_def, ws),
axis(OperatorBase::GetSingleArg<int>("axis", -1)),
momentum(OperatorBase::GetSingleArg<float>("momentum", float(0.9))),
eps(OperatorBase::GetSingleArg<float>("eps", float(1e-3))),
momentum(OperatorBase::GetSingleArg<float>("momentum", 0.9f)),
eps(OperatorBase::GetSingleArg<float>("eps", 1e-3f)),
use_stats(OperatorBase::GetSingleArg<int>("use_stats", -1)) {}
USE_OPERATOR_FUNCTIONS(Context);
......@@ -112,7 +114,7 @@ class FusedBatchNormGradientOp : public Operator<Context> {
FusedBatchNormGradientOp(const OperatorDef& op_def, Workspace* ws)
: Operator<Context>(op_def, ws),
axis(OperatorBase::GetSingleArg<int>("axis", -1)),
eps(OperatorBase::GetSingleArg<float>("eps", float(1e-3))),
eps(OperatorBase::GetSingleArg<float>("eps", 1e-3f)),
use_stats(OperatorBase::GetSingleArg<int>("use_stats", -1)) {}
USE_OPERATOR_FUNCTIONS(Context);
......@@ -143,11 +145,16 @@ template <class Context>
class CuDNNBatchNormOp final : public FusedBatchNormOp<Context> {
public:
CuDNNBatchNormOp(const OperatorDef& op_def, Workspace* ws)
: FusedBatchNormOp<Context>(op_def, ws) {
: FusedBatchNormOp<Context>(op_def, ws),
eps64(OperatorBase::GetSingleArg<float>("eps", 1e-3f)) {
CUDNN_CHECK(cudnnCreateTensorDescriptor(&input_desc));
CUDNN_CHECK(cudnnCreateTensorDescriptor(&output_desc));
CUDNN_CHECK(cudnnCreateTensorDescriptor(&bn_desc));
this->eps = std::max(this->eps, float(CUDNN_BN_MIN_EPSILON));
if (eps64 <= CUDNN_BN_MIN_EPSILON - FLT_EPSILON)
LOG(FATAL) << "Provided epsilon is smaller than "
<< "CUDNN_BN_MIN_EPSILON. Setting it to "
<< "CUDNN_BN_MIN_EPSILON instead.";
eps64 = std::max(eps64, CUDNN_BN_MIN_EPSILON);
}
USE_OPERATOR_FUNCTIONS(Context);
......@@ -163,6 +170,7 @@ class CuDNNBatchNormOp final : public FusedBatchNormOp<Context> {
template <typename T> void RunWithType();
protected:
double eps64;
cudnnTensorDescriptor_t input_desc, output_desc, bn_desc;
cudnnBatchNormMode_t bn_mode;
TIndex N, C;
......@@ -174,11 +182,16 @@ template <class Context>
class CuDNNBatchNormGradientOp final : public FusedBatchNormGradientOp<Context> {
public:
CuDNNBatchNormGradientOp(const OperatorDef& op_def, Workspace* ws)
: FusedBatchNormGradientOp<Context>(op_def, ws) {
: FusedBatchNormGradientOp<Context>(op_def, ws),
eps64(OperatorBase::GetSingleArg<float>("eps", 1e-3f)) {
CUDNN_CHECK(cudnnCreateTensorDescriptor(&input_desc));
CUDNN_CHECK(cudnnCreateTensorDescriptor(&output_desc));
CUDNN_CHECK(cudnnCreateTensorDescriptor(&bn_desc));
this->eps = std::max(this->eps, float(CUDNN_BN_MIN_EPSILON));
if (eps64 <= CUDNN_BN_MIN_EPSILON - FLT_EPSILON)
LOG(FATAL) << "Provided epsilon is smaller than "
<< "CUDNN_BN_MIN_EPSILON. Setting it to "
<< "CUDNN_BN_MIN_EPSILON instead.";
eps64 = std::max(eps64, CUDNN_BN_MIN_EPSILON);
}
USE_OPERATOR_FUNCTIONS(Context);
......@@ -195,6 +208,7 @@ class CuDNNBatchNormGradientOp final : public FusedBatchNormGradientOp<Context>
template <typename T> void InferenceRunWithType();
protected:
double eps64;
cudnnTensorDescriptor_t input_desc, output_desc, bn_desc;
cudnnBatchNormMode_t bn_mode;
TIndex N, C, S, NC, NS;
......
Dragon/include/operators/norm/batch_renorm_op.h
......@@ -22,13 +22,13 @@ class BatchRenormOp : public Operator<Context> {
BatchRenormOp(const OperatorDef& op_def, Workspace* ws)
: Operator<Context>(op_def, ws),
axis(OperatorBase::GetSingleArg<int>("axis", -1)),
momentum(OperatorBase::GetSingleArg<float>("momentum", float(0.9))),
eps(OperatorBase::GetSingleArg<float>("eps", float(1e-3))),
r_max(OperatorBase::GetSingleArg<float>("r_max", float(3.0))),
d_max(OperatorBase::GetSingleArg<float>("d_max", float(5.0))),
t_delta(OperatorBase::GetSingleArg<float>("t_delta", float(1.0))),
momentum(OperatorBase::GetSingleArg<float>("momentum", 0.9f)),
eps(OperatorBase::GetSingleArg<float>("eps", 1e-3f)),
r_max(OperatorBase::GetSingleArg<float>("r_max", 3.f)),
d_max(OperatorBase::GetSingleArg<float>("d_max", 5.f)),
t_delta(OperatorBase::GetSingleArg<float>("t_delta", 1.f)),
use_stats(OperatorBase::GetSingleArg<int>("use_stats", -1)),
t_r_max(float(1.0)), t_d_max(float(0.0)), t_val(float(0.0)),
t_r_max(1.f), t_d_max(0.f), t_val(0.f),
mode(OperatorBase::GetSingleArg<string>("mode", "DEFAULT")) {
if (axis != -1)
CHECK_EQ(axis, 1)
......
Dragon/include/operators/norm/group_norm_op.h
......@@ -23,7 +23,7 @@ class GroupNormOp : public Operator<Context> {
: Operator<Context>(op_def, ws),
group(OperatorBase::GetSingleArg<int>("group", 32)),
axis(OperatorBase::GetSingleArg<int>("axis", -1)),
eps(OperatorBase::GetSingleArg<float>("eps", float(1e-3))) {
eps(OperatorBase::GetSingleArg<float>("eps", 1e-3f)) {
if (axis != -1)
CHECK_EQ(axis, 1)
<< "\nThe axis can only be set to 1.";
......@@ -77,7 +77,7 @@ class FusedGroupNormOp : public Operator<Context> {
: Operator<Context>(op_def, ws),
group(OperatorBase::GetSingleArg<int>("group", 32)),
axis(OperatorBase::GetSingleArg<int>("axis", -1)),
eps(OperatorBase::GetSingleArg<float>("eps", float(1e-3))) {}
eps(OperatorBase::GetSingleArg<float>("eps", 1e-3f)) {}
USE_OPERATOR_FUNCTIONS(Context);
void Setup();
......@@ -100,8 +100,7 @@ class FusedGroupNormGradientOp : public Operator<Context> {
FusedGroupNormGradientOp(const OperatorDef& op_def, Workspace* ws)
: Operator<Context>(op_def, ws),
group(OperatorBase::GetSingleArg<int>("group", 32)),
axis(OperatorBase::GetSingleArg<int>("axis", -1)),
eps(OperatorBase::GetSingleArg<float>("eps", float(1e-3))) {}
axis(OperatorBase::GetSingleArg<int>("axis", -1)) {}
USE_OPERATOR_FUNCTIONS(Context);
void Setup();
......@@ -110,7 +109,6 @@ class FusedGroupNormGradientOp : public Operator<Context> {
template <typename T> void RunWithType();
protected:
float eps;
Tensor num_by_chans;
Tensor* multiplier, *num_multiplier, *spatial_multiplier, *cgs_multiplier;
Tensor* mean, *var, *stddev, *x_norm;
......
Dragon/include/operators/norm/instance_norm_op.h
......@@ -22,7 +22,7 @@ class InstanceNormOp : public Operator<Context> {
InstanceNormOp(const OperatorDef& op_def, Workspace* ws)
: Operator<Context>(op_def, ws),
axis(OperatorBase::GetSingleArg<int>("axis", -1)),
eps(OperatorBase::GetSingleArg<float>("eps", float(1e-3))) {
eps(OperatorBase::GetSingleArg<float>("eps", 1e-3f)) {
if (axis != -1)
CHECK_EQ(axis, 1)
<< "\nThe axis can only be set to 1.";
......
Dragon/include/operators/norm/l2_norm_op.h
......@@ -23,7 +23,7 @@ class L2NormOp final : public Operator<Context> {
: Operator<Context>(op_def, ws),
axis(OperatorBase::GetSingleArg<int>("axis", 0)),
num_axes(OperatorBase::GetSingleArg<int>("num_axes", -1)),
eps(OperatorBase::GetSingleArg<float>("eps", float(1e-5))),
eps(OperatorBase::GetSingleArg<float>("eps", 1e-5f)),
mode(OperatorBase::GetSingleArg<string>("mode", "SUM")) {}
USE_OPERATOR_FUNCTIONS(Context);
......
Dragon/include/operators/update/sgd_update_op.h
......@@ -20,14 +20,15 @@ template <class Context>
class SGDUpdateOp final : public UpdateOpBase<Context> {
public:
SGDUpdateOp(const OperatorDef& op_def, Workspace* ws)
: UpdateOpBase<Context>(op_def, ws) {}
: UpdateOpBase<Context>(op_def, ws),
old_lr(-1.f), correction(1.f) {}
USE_OPERATOR_FUNCTIONS(Context);
USE_UPDATER_FUNCTIONS(Context);
void ComputeRunWithFloat() override;
protected:
float lr, momentum;
float old_lr, lr, momentum, correction;
};
} // namespace dragon
......
Dragon/include/utils/op_kernel.h
......@@ -286,6 +286,11 @@ void SparseSoftmaxFocalLossGrad(const int count,
Tensor* ignore,
T* dx);
/******************** misc.dtype ********************/
template <typename Ta, typename Tb, class Context>
void TypeA2B(const int count, const Ta* a, Tb* b);
/******************** misc.image_data ********************/
template <typename Tx, typename Ty, class Context>
......@@ -308,25 +313,25 @@ void Arange(const int count,
const int step,
T* y);
/******************** ndarray.argmax ********************/
/******************** ndarray.argreduce ********************/
template <typename T, class Context>
void Argmax(const int count,
void Argmax(const int count,
const int axis_dim,
const int inner_dim,
const int top_k,
const T* x,
T* y);
/******************** ndarray.argmin ********************/
const int inner_dim,
const int top_k,
const T* x,
int64_t* indices,
T* values);
template <typename T, class Context>
void Argmin(const int count,
void Argmin(const int count,
const int axis_dim,
const int inner_dim,
const int top_k,
const T* x,
T* y);
const int inner_dim,
const int top_k,
const T* x,
int64_t* indices,
T* values);
/******************** ndarray.gather ********************/
......
Dragon/python/dragon/core/gradient_maker.py
......@@ -181,7 +181,7 @@ class GraphGradientMaker(object):
# Append ops
if not is_skip:
# + Gen Op
# --> GenOp
if len(gen_grads) > 0:
op_inputs = []; op_outputs = []; values = []
for item in gen_grads:
......@@ -193,7 +193,7 @@ class GraphGradientMaker(object):
if forward_op.HasField('device_option'):
gen_op.device_option.CopyFrom(forward_op.device_option)
backward_ops.append(gen_op)
# + Grad Op
# --> GradOp
for g_op in g_ops:
g_op.name = _op_name(None if auto_names else 'runtime')
backward_ops.append(g_op)
......@@ -204,8 +204,10 @@ class GraphGradientMaker(object):
original_idx = -1
for g_input_idx, g_input in enumerate(g_inputs):
if g_output == g_input: original_idx = g_input_idx
# Ignore un-used GI(?)
# Ignore un-used && in-placed GI(?)
if original_idx == -1: continue
if g_output in g_op.input: continue
# Found a split branch
original_name = forward_op.input[original_idx]
if inputs_count[original_name] > 1:
# Split
......
Dragon/python/dragon/core/tensor.py
......@@ -340,6 +340,34 @@ class Tensor(object):
def dtype(self, value):
self._dtype = value
def astype(self, dtype, inplace=False):
"""Cast the data type of inputs to a specific one.
Parameters
----------
dtype : str
The specific dtype.
inplace : boolean
Whether to modify the inputs.
Returns
-------
Tensor
The output tensor.
"""
if inplace:
output = Tensor.CreateOperator(inputs=[], existing_outputs=[self],
op_type='AsType', dtype=dtype)
else:
output = Tensor.CreateOperator(inputs=self, nout=1,
op_type='AsType', dtype=dtype)
if self.shape is not None:
output.shape = self.shape[:]
return output
@property
def extra_targets(self):
"""Return or Set the extra solving targets.
......@@ -646,6 +674,8 @@ class Tensor(object):
output.shape = self.shape[:]
return output
__truediv__ = __div__
def __rdiv__(self, other):
"""Calculate y / x.
......@@ -672,6 +702,8 @@ class Tensor(object):
output.shape = self.shape[:]
return output
__rtruediv__ = __rdiv__
def __neg__(self):
"""Calculate -x.
......
Dragon/python/dragon/docs/contents/operators/cast.rst deleted 100644 → 0
===========
:mod:`Cast`
===========
.. toctree::
:hidden:
.. automodule:: dragon.operators.cast
:members:
\ No newline at end of file
Dragon/python/dragon/docs/contents/ops.rst
......@@ -129,7 +129,9 @@ List Brief
`Reduce`_ The general reduce operator.
`Sum`_ Compute the sum along the given axis.
`Mean`_ Compute the mean along the given axis.
`Max`_ Compute the values of maximum elements along the given axis.
`Argmax`_ Compute the indices of maximum elements along the given axis.
`Min`_ Compute the values of minimum elements along the given axis.
`Argmin`_ Compute the indices of minimum elements along the given axis.
`Slice`_ Slice interface of NDArray.
`Stack`_ Stack the inputs along the given axis.
......@@ -161,6 +163,7 @@ Misc
================= ======================================================================
List Brief
================= ======================================================================
`AsType`_ Cast the data type of inputs to a specific one.
`Run`_ Run a custom operator. (Without GradientFlow)
`Template`_ Run a custom operator. (With GradientFlow)
`Accuracy`_ Calculate the Top-K accuracy.
......@@ -176,14 +179,6 @@ List Brief
`Proposal`_ Generate Regional Proposals, introduced by `[Ren et.al, 2015] <https://arxiv.org/abs/1506.01497>`_.
================= ======================================================================
Cast
----
================= ======================================================================
List Brief
================= ======================================================================
`FloatToHalf`_ Cast the type of tensor from ``float32`` to ``float16``.
================= ======================================================================
MPI
---
================= ======================================================================
......@@ -266,7 +261,9 @@ List Brief
.. _Reduce: operators/ndarray.html#dragon.operators.ndarray.Reduce
.. _Sum: operators/ndarray.html#dragon.operators.ndarray.Sum
.. _Mean: operators/ndarray.html#dragon.operators.ndarray.Mean
.. _Max: operators/ndarray.html#dragon.operators.ndarray.Max
.. _Argmax: operators/ndarray.html#dragon.operators.ndarray.Argmax
.. _Min: operators/ndarray.html#dragon.operators.ndarray.Min
.. _Argmin: operators/ndarray.html#dragon.operators.ndarray.Argmin
.. _Slice: operators/ndarray.html#dragon.operators.ndarray.Slice
.. _Stack: operators/ndarray.html#dragon.operators.ndarray.Stack
......@@ -285,6 +282,7 @@ List Brief
.. _Copy: operators/control_flow.html#dragon.operators.control_flow.Copy
.. _Equal: operators/control_flow.html#dragon.operators.control_flow.Equal
.. _AsType: operators/misc.html#dragon.operators.misc.AsType
.. _Run: operators/misc.html#dragon.operators.misc.Run
.. _Template: operators/misc.html#dragon.operators.misc.Template
.. _Accuracy: operators/misc.html#dragon.operators.misc.Accuracy
......@@ -293,7 +291,5 @@ List Brief
.. _Proposal: operators/contrib/rcnn.html#dragon.operators.contrib.rcnn.ops.Proposal
.. _FloatToHalf: operators/cast.html#dragon.operators.misc.FloatToHalf
.. _MPIBroadcast: operators/mpi.html#dragon.operators.mpi.MPIBroadcast
.. _MPIGather: operators/mpi.html#dragon.operators.mpi.MPIGather
\ No newline at end of file
Dragon/python/dragon/io/blob_fetcher.py
......@@ -62,7 +62,7 @@ class BlobFetcher(Process):
# fill blobs
im, labels = self.Q_in.get()
im_blob = np.zeros(shape=([self._batch_size] + list(im.shape)), dtype=np.uint8)
label_blob = np.zeros((self._batch_size, len(labels)), dtype=np.float32)
label_blob = np.zeros((self._batch_size, len(labels)), dtype=np.int64)
for ix in range(0, self._batch_size):
im_blob[ix, :, :, :], label_blob[ix, :] = im, labels
if ix != self._batch_size - 1: im, labels = self.Q_in.get()
......
Dragon/python/dragon/io/data_batch.py
......@@ -38,10 +38,10 @@ class DataBatch(object):
----------
source : str
The path of database.
multiple_nodes: boolean
Whether to split data for multiple parallel nodes.
shuffle : boolean
Whether to shuffle the data.
node_step: boolean
Whether to split data for multiple parallel nodes.
num_chunks : int
The number of chunks to split. Default is ``2048``.
chunk_size : int
......@@ -103,8 +103,12 @@ class DataBatch(object):
self._num_transformers += 1
# add 1 transformer for random scale
if kwargs.get('max_random_scale', 1.0) - \
kwargs.get('min_random_scale', 1.0) != 0:
self._num_transformers += 1
kwargs.get('min_random_scale', 1.0) != 0:
self._num_transformers += 1
# add 1 transformer for random crop
if kwargs.get('crop_size', 0) > 0 and \
kwargs.get('phase', 'TEST') == 'TRAIN':
self._num_transformers += 1
self._num_transformers = min(self._num_transformers, self._max_transformers)
self._batch_size = kwargs.get('batch_size', 100)
......@@ -127,8 +131,8 @@ class DataBatch(object):
num_parts = self._num_readers
part_idx = i
if self._readers[i]._use_shuffle \
or self._readers[i]._use_step:
if self._readers[i]._multiple_nodes or \
self._readers[i]._use_shuffle:
num_parts *= group_size
part_idx += local_rank * self._num_readers
......
Dragon/python/dragon/io/data_reader.py
......@@ -14,6 +14,7 @@ from __future__ import division
from __future__ import print_function
import math
import numpy as np
import numpy.random as npr
from multiprocessing import Process
......@@ -34,10 +35,12 @@ class DataReader(Process):
----------
source : str
The path of database.
multiple_nodes: boolean
Whether to split data for multiple parallel nodes.
shuffle : boolean
Whether to shuffle the data.
node_step: boolean
Whether to split data for multiple parallel nodes.
instance_chunk : boolean
Whether to limit each chunk with at most 1 instance.
num_chunks : int
The number of chunks to split. Default is ``2048``.
chunk_size : int
......@@ -46,8 +49,9 @@ class DataReader(Process):
"""
super(DataReader, self).__init__()
self._source = kwargs.get('source', '')
self._multiple_nodes = kwargs.get('multiple_nodes', False)
self._use_shuffle = kwargs.get('shuffle', False)
self._use_step = kwargs.get('node_step', False)
self._use_instance_chunk = kwargs.get('instance_chunk', False)
self._num_chunks = kwargs.get('num_chunks', 2048)
self._chunk_size = kwargs.get('chunk_size', -1)
......@@ -106,9 +110,8 @@ class DataReader(Process):
None
"""
if self._use_shuffle or self._use_step:
if self._use_shuffle:
self._perm = npr.permutation(self._num_shuffle_parts)
if self._multiple_nodes or self._use_shuffle:
if self._use_shuffle: self._perm = npr.permutation(self._num_shuffle_parts)
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._start_idx * self._chunk_size)
......@@ -168,16 +171,21 @@ class DataReader(Process):
self._db_size = int(self._db.get('size'))
self._db_zfill = int(self._db.get('zfill'))
self._epoch_size = int(self._db_size / self._num_parts + 1)
# search a optimal chunk size by chunks
if self._chunk_size == -1:
max_chunk_size = self._db._total_size / ((self._num_chunks * (1 << 20)))
min_chunk_size = 1
while min_chunk_size * 2 < max_chunk_size: min_chunk_size *= 2
self._chunk_size = min_chunk_size
self._num_shuffle_parts = int(math.ceil(self._db._total_size * 1.1 /
(self._num_parts * self._chunk_size << 20)))
self._chunk_size = int(self._db_size / self._num_shuffle_parts / self._num_parts + 1)
self._perm = npr.permutation(self._num_shuffle_parts)
if self._use_instance_chunk:
self._chunk_size = 1
self._num_shuffle_parts = int(self._db_size / self._chunk_size / self._num_parts) + 1
else:
# search a optimal chunk size by chunks
if self._chunk_size == -1:
max_chunk_size = self._db._total_size / ((self._num_chunks * (1 << 20)))
min_chunk_size = 1
while min_chunk_size * 2 < max_chunk_size: min_chunk_size *= 2
self._chunk_size = min_chunk_size
self._num_shuffle_parts = int(math.ceil(self._db._total_size * 1.1 /
(self._num_parts * self._chunk_size << 20)))
self._chunk_size = int(self._db_size / self._num_shuffle_parts / self._num_parts + 1)
self._perm = np.arange(self._num_shuffle_parts)
# init env
self.reset()
......@@ -187,5 +195,6 @@ class DataReader(Process):
self.Q_out.put(self.element())
self.next_record()
if self._cur_idx >= self._end_idx:
if self._use_shuffle or self._use_step: self.next_chunk()
if self._multiple_nodes or \
self._use_shuffle: self.next_chunk()
else: self.reset()
\ No newline at end of file
Dragon/python/dragon/io/data_transformer.py
......@@ -13,6 +13,7 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import sys
import numpy as np
import numpy.random as npr
from multiprocessing import Process
......@@ -28,7 +29,7 @@ try:
import PIL.Image
import PIL.ImageEnhance
except ImportError as e:
print('Failed to import PIL. Error: {0}'.format(str(e)))
print("Failed to import PIL. \nIt's OK if disabling color augmentation.".format(str(e)))
class DataTransformer(Process):
......@@ -103,12 +104,20 @@ class DataTransformer(Process):
random_scale = npr.uniform() * (self._max_random_scale - self._min_random_scale) \
+ self._min_random_scale
if random_scale != 1.0:
im = cv2.resize(im, None, interpolation=cv2.INTER_LINEAR,
fx=random_scale, fy=random_scale)
if sys.version_info >= (3, 0):
im = cv2.resize(im, None, interpolation=cv2.INTER_LINEAR,
fx=random_scale, fy=random_scale)
else:
# Fuck Fuck Fuck opencv-python2, it always has a BUG
# that leads to duplicate cuDA handles created at gpu:0
new_shape = (int(im.shape[1] * random_scale), int(im.shape[0] * random_scale))
im = PIL.Image.fromarray(im)
im = im.resize(new_shape, PIL.Image.BILINEAR)
im = np.array(im)
# random crop
if self._crop_size > 0:
if self._phase== 'TRAIN':
if self._phase == 'TRAIN':
h_off = npr.randint(im.shape[0] - self._crop_size + 1)
w_off = npr.randint(im.shape[1] - self._crop_size + 1)
else:
......@@ -152,7 +161,12 @@ class DataTransformer(Process):
self._padding : self._padding + im.shape[1], :] = im
im = pad_img
return im, [datum.label]
# labels
labels = []
if len(datum.labels) > 0: labels.extend(datum.labels)
else: labels.append(datum.label)
return im, labels
def run(self):
"""Start the process.
......
Dragon/python/dragon/operators/activation.py
......@@ -200,7 +200,9 @@ def Tanh(inputs, **kwargs):
def Dropout(inputs, prob=0.5, scale=True, **kwargs):
"""Randomly set a unit into zero, introduced by `[Srivastava et.al, 2014] <http://jmlr.org/papers/v15/srivastava14a.html>`_.
"""Randomly set a unit into zero.
Introduced by `[Srivastava et.al, 2014] <http://jmlr.org/papers/v15/srivastava14a.html>`_.
Parameters
----------
......
Dragon/python/dragon/operators/cast.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>
#
# ------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from . import *
def FloatToHalf(inputs, **kwargs):
"""Cast the type of tensor from ``float32`` to ``float16``.
Parameters
----------
inputs : Tensor
The ``float32`` tensor.
Returns
-------
Tensor
The ``float16`` tensor.
"""
CheckInputs(inputs, 1)
arguments = ParseArguments(locals())
output = Tensor.CreateOperator(nout=1, op_type='FloatToHalf', **arguments)
if inputs.shape is not None:
output.shape = inputs.shape[:]
return output
\ No newline at end of file
Dragon/python/dragon/operators/misc.py
......@@ -16,6 +16,49 @@ from __future__ import print_function
from . import *
def AsType(inputs, dtype='float32', inplace=False, **kwargs):
"""Cast the data type of inputs to a specific one.
If ``inplace`` is ``True``, cast ``self`` instead of returning a new one.
Parameters
----------
inputs : Tensor
The input tensor.
dtype : str
The specific data type.
inplace : boolean
Whether to modify the inputs.
Returns
-------
Tensor
The output tensor.
Examples
--------
>>> x = Tensor('x', dtype='float32').Variable()
>>> y = AsType(x, 'int32')
>>> z = x.astype('int64')
>>> xx = x.astype('float64', inplace=True)
>>> print(x.name, xx.name)
"""
CheckInputs(inputs, 1)
arguments = ParseArguments(locals())
if inplace:
arguments['inputs'] = []
arguments['existing_outputs'] = [inputs]
output = Tensor.CreateOperator(nout=1, op_type='AsType', **arguments)
if inputs.shape is not None:
output.shape = inputs.shape[:]
return output
def Run(inputs, module, op, param_str='', nout=1, **kwargs):
"""Run a custom operator. (Without GradientFlow)
......
Dragon/python/dragon/operators/ndarray.py
......@@ -235,7 +235,6 @@ def Concat(inputs, axis=1, **kwargs):
"""
CheckInputs(inputs, 1, INT_MAX)
arguments = ParseArguments(locals())
arguments['num_input'] = len(inputs)
output = Tensor.CreateOperator(nout=1, op_type='Concat', **arguments)
......@@ -337,6 +336,26 @@ def Mean(inputs, axis=-1, keep_dims=False, **kwargs):
return Reduce(inputs, axis, 'MEAN', keep_dims, **kwargs)
def _ArgReduce(inputs, axis=-1, operation='NONE', top_k=1, keep_dims=False, **kwargs):
CheckInputs(inputs, 1)
arguments = ParseArguments(locals())
n_out = 1
if 'ARG' not in operation:
n_out = 2; arguments['operation'] = 'ARG' + operation
outputs = Tensor.CreateOperator(nout=n_out, op_type='ArgReduce', **arguments)
if 'ARG' not in operation: output = outputs[1]
else: output = outputs
if inputs.shape is not None:
output.shape = inputs.shape[:]
if top_k > 1: output.shape[axis] = top_k
else: del output.shape[axis]
return output
def Argmax(inputs, axis=-1, top_k=1, keep_dims=False, **kwargs):
"""Compute the indices of maximum elements along the given axis.
......@@ -357,17 +376,30 @@ def Argmax(inputs, axis=-1, top_k=1, keep_dims=False, **kwargs):
The indices.
"""
CheckInputs(inputs, 1)
arguments = ParseArguments(locals())
return _ArgReduce(inputs, axis, 'ARGMAX', top_k, keep_dims, **kwargs)
output = Tensor.CreateOperator(nout=1, op_type='Argmax', **arguments)
if inputs.shape is not None:
output.shape = inputs.shape[:]
if top_k > 1: output.shape[axis] = top_k
else: del output.shape[axis]
def Max(inputs, axis=-1, top_k=1, keep_dims=False, **kwargs):
"""Compute the values of maximum elements along the given axis.
return output
Parameters
----------
inputs : Tensor
The input tensor.
axis : int
The axis to compute. Default is ``-1`` (Along all axes).
top_k : int
The top k results to keep.
keep_dims : boolean
Whether to keep dims after computing.
Returns
-------
Tensor
The values.
"""
return _ArgReduce(inputs, axis, 'MAX', top_k, keep_dims, **kwargs)
def Argmin(inputs, axis=-1, top_k=1, keep_dims=False, **kwargs):
......@@ -390,17 +422,30 @@ def Argmin(inputs, axis=-1, top_k=1, keep_dims=False, **kwargs):
The indices.
"""
CheckInputs(inputs, 1)
arguments = ParseArguments(locals())
return _ArgReduce(inputs, axis, 'ARGMIN', top_k, keep_dims, **kwargs)
output = Tensor.CreateOperator(nout=1, op_type='Argmin', **arguments)
if inputs.shape is not None:
output.shape = inputs.shape[:]
if top_k > 1: output.shape[axis] = top_k
else: del output.shape[axis]
def Min(inputs, axis=-1, top_k=1, keep_dims=False, **kwargs):
"""Compute the values of minimum elements along the given axis.
return output
Parameters
----------
inputs : Tensor
The input tensor.
axis : int
The axis to compute. Default is ``-1`` (Along all axes).
top_k : int
The top k results to keep.
keep_dims : boolean
Whether to keep dims after computing.
Returns
-------
Tensor
The values.
"""
return _ArgReduce(inputs, axis, 'MIN', top_k, keep_dims, **kwargs)
def Transpose(inputs, perms=None, **kwargs):
......
Dragon/python/dragon/ops.py
......@@ -21,7 +21,6 @@ from .operators import activation as act
from .operators import arithmetic as math
from .operators import control_flow
from .operators import misc as misc
from .operators import cast
from .operators import mpi
from .operators import ndarray
from .operators import norm
......@@ -110,7 +109,9 @@ Crop = ndarray.Crop
Reduce = ndarray.Reduce
Sum = ndarray.Sum
Mean = ndarray.Mean
Max = ndarray.Max
Argmax = ndarray.Argmax
Min = ndarray.Min
Argmin = ndarray.Argmin
Slice = ndarray.Slice
Stack = ndarray.Stack
......@@ -131,15 +132,13 @@ Copy = control_flow.Copy
Equal = control_flow.Equal
# misc
AsType = misc.AsType
Run = misc.Run
Template = misc.Template
Accuracy = misc.Accuracy
StopGradient = misc.StopGradient
MovingAverage = misc.MovingAverage
# cast
FloatToHalf = cast.FloatToHalf
# mpi
MPIBroadcast = mpi.MPIBroadcast
MPIGather = mpi.MPIGather
......
Dragon/python/dragon/vm/caffe/proto/caffe.proto
......@@ -38,6 +38,7 @@ message Datum {
repeated float float_data = 6;
// If true data contains an encoded image that need to be decoded
optional bool encoded = 7 [default = false];
repeated int32 labels = 8;
}
message FillerParameter {
......@@ -1459,8 +1460,8 @@ message NormalizeParameter {
}
message ParallelParameter {
optional bool shuffle = 1 [default = false];
optional bool node_step = 2 [default = false];
optional bool multiple_nodes = 1 [default = false];
optional bool shuffle = 2 [default = false];
optional bool partition = 3 [default = false];
}
......
Dragon/python/dragon/vm/torch/execute_engine.py
......@@ -36,6 +36,8 @@ from dragon.config import option
from .autograd.expression import Expression
from .autograd.grad_mode import is_grad_enabled
from .constants import CTX_TO_DEVICE_OPTION
from .tensor import RuntimeTensor
from .tensor_pool import TPool
def RunOperator(inputs, outputs, meta, auto_grad=True, **kwargs):
......@@ -44,36 +46,41 @@ def RunOperator(inputs, outputs, meta, auto_grad=True, **kwargs):
if len(outputs) == 0:
raise ValueError('The num of outputs should be at least 1.')
# + I/O Check
requires_grad = False
inputs_name = []; outputs_name = []
for input in inputs:
inputs_name.append(input.name)
if input.requires_grad: requires_grad = True
requires_grad = requires_grad and is_grad_enabled()
inputs_name = [input.name for input in inputs]
outputs_name = [output.name for output in outputs]
for ix, output in enumerate(outputs):
if isinstance(output, tuple):
name = TPool.get('join' if requires_grad else 'detach')
outputs[ix] = RuntimeTensor(name, dtype=output[0], ctx=output[1])
outputs_name.append(outputs[ix].name)
# + Engine Check
engine_type = meta[0]; persistent_key = None
if engine_type == 'ONCE':
# OpType + CTX -> Op
# ++ OpType + CTX -> Op
op_type, ctx = meta[1:]
if ctx is None: raise ValueError('Excepted a context, got None.')
op = pb.MakeOperatorDef(op_type, inputs_name, outputs_name, name='runtime',
device_option=CTX_TO_DEVICE_OPTION[ctx], **kwargs)
device_option=CTX_TO_DEVICE_OPTION[ctx], **kwargs)
elif engine_type == 'PERSISTENT':
# Key + Inputs + Outputs -> Op
# ++ Key + Inputs + Outputs -> Op
persistent_key, meta_op = meta[1:]
op = pb.MutableOperatorDef(meta_op, inputs_name, outputs_name)
else:
raise ValueError('Unknown executing engine: {}.'.format(engine_type))
# ~ Auto-Grad ~
requires_grad = False
for input in inputs:
if input.requires_grad:
requires_grad = True
# + Auto-Grad
if len(inputs) > 0 and auto_grad:
input_expressions = []
if requires_grad and is_grad_enabled():
if requires_grad:
ignored_grads = set()
# + Trace outputs
# ++ Trace outputs
for input in inputs:
input_expressions.append(input._expr)
if input._ignored_grads:
......@@ -88,11 +95,11 @@ def RunOperator(inputs, outputs, meta, auto_grad=True, **kwargs):
if len(ignored_grads) > 0:
outputs[ix]._ignored_grads = ignored_grads
else:
# + Reset status
# ++ Reset status
for ix in range(len(outputs)):
outputs[ix]._requires_grad = False
# Run
# + Run
if option['log_optimized_graph'] or option['log_meta_graph']:
from dragon.config import logger
logger.info('>>>>>>>>>>>>>>>>>> Forward Flow <<<<<<<<<<<<<<<<<<\n')
......@@ -104,7 +111,7 @@ def RunOperator(inputs, outputs, meta, auto_grad=True, **kwargs):
dg.workspace.RunPersistentOp(persistent_key,
op.name, inputs_name, outputs_name)
# Returns
# + Returns
if len(outputs) > 1: return outputs
elif len(outputs) == 1: return outputs[0]
else: return None
\ No newline at end of file
Dragon/python/dragon/vm/torch/module.py
......@@ -30,14 +30,14 @@ from dragon.config import logger
from dragon.vm.torch.environ import \
add_submodule, get_module_name
from dragon.vm.torch.tensor import Tensor, RuntimeTensor, Parameter
from dragon.vm.torch.tensor import Tensor, Parameter
from dragon.vm.torch.tensor_pool import TPool
from dragon.vm.torch.execute_engine import RunOperator
class Module(object):
def __init__(self):
self._modules = {}
self._modules = OrderedDict()
self._parameters = OrderedDict()
self._buffers = OrderedDict()
self._persistent_key = self._op = None
......@@ -96,22 +96,24 @@ class Module(object):
else:
object.__setattr__(self, key, value)
def state_dict(self, destination=None, prefix='', keep_vars=False):
def state_dict(self, destination=None, prefix='', to_numpy=False):
if destination is None:
destination = OrderedDict()
for name, param in self._parameters.items():
if param is not None:
destination[prefix + name] = param if keep_vars else param.data
destination[prefix + name] = \
param.cpu().numpy().copy() if to_numpy else param
for name, buf in self._buffers.items():
if buf is not None:
destination[prefix + name] = buf
destination[prefix + name] = \
buf.cpu().numpy().copy() if to_numpy else buf
for name, module in self._modules.items():
if module is not None:
module.state_dict(destination, prefix + name + '.', keep_vars=keep_vars)
module.state_dict(destination, prefix + name + '.', to_numpy=to_numpy)
return destination
def load_state_dict(self, state_dict, strict=True):
logger.info('Load state dict from the stream.')
logger.info('Load the state dict from numpy arrays.')
def submodule_key_mismatch(full_name, is_missing):
module = self
names = full_name.split(".")
......@@ -198,6 +200,14 @@ class Module(object):
else:
self._parameters[name] = param
def add_module(self, name, module, name_v2=None):
if not isinstance(module, Module) and module is not None:
raise TypeError("{} is not a Module subclass".format(type(module)))
if hasattr(self, name) and name not in self._modules:
raise KeyError("attribute '{}' already exists".format(name))
self._modules[name] = module
add_submodule(module, name_v2 if name_v2 else name)
def __call__(self, *args, **kwargs):
with dg.name_scope(get_module_name(self)):
return self.forward(*args)
......@@ -238,6 +248,23 @@ class Module(object):
memo.add(module)
yield name, module
def modules(self):
for name, module in self.named_modules():
yield module
def named_modules(self, memo=None, prefix=''):
if memo is None:
memo = set()
if self not in memo:
memo.add(self)
yield prefix, self
for name, module in self._modules.items():
if module is None:
continue
submodule_prefix = prefix + ('.' if prefix else '') + name
for m in module.named_modules(memo, submodule_prefix):
yield m
def named_parameters(self, memo=None, prefix=''):
"""Returns an iterator over module parameters.
......@@ -315,7 +342,7 @@ class Module(object):
raise NotImplementedError()
def register_output(self, dtype='float32'):
return RuntimeTensor(dtype=dtype, ctx=self._ctx)
return (dtype, self._ctx)
def unify_devices(self, inputs):
for ix, t in enumerate(inputs):
......
Dragon/python/dragon/vm/torch/nn/__init__.py
......@@ -18,8 +18,10 @@ from dragon.vm.torch.module import Module
from dragon.vm.torch.tensor import Parameter
from .modules.conv import Conv2d
from .modules.pooling import MaxPool2d, AvgPool2d
from .modules.activation import ReLU
from .modules.activation import ReLU, Softmax
from .modules.linear import Linear
from .modules.loss import CrossEntropyLoss
from .modules.container import Container, Sequential, ModuleList
from .modules.batchnorm import BatchNorm1d, BatchNorm2d, BatchNorm3d
from .modules.dropout import Dropout, Dropout2d, Dropout3d
from . import init
\ No newline at end of file
Dragon/python/dragon/vm/torch/nn/init.py
......@@ -19,6 +19,7 @@ from __future__ import print_function
import math
import warnings
from dragon.vm.torch.autograd.grad_mode import no_grad
def calculate_gain(nonlinearity, param=None):
......@@ -43,15 +44,18 @@ def calculate_gain(nonlinearity, param=None):
def uniform_(tensor, a=0, b=1):
return tensor.uniform_(a, b)
with no_grad():
return tensor.uniform_(a, b)
def normal_(tensor, mean=0, std=1):
return tensor.normal_(mean, std)
with no_grad():
return tensor.normal_(mean, std)
def constant_(tensor, val):
return tensor.fill_(val)
with no_grad():
return tensor.fill_(val)
def dirac_(tensor):
......@@ -61,16 +65,16 @@ def dirac_(tensor):
sizes = tensor.size()
min_dim = min(sizes[0], sizes[1])
tensor.zero_()
for d in range(min_dim):
if dimensions == 3: # Temporal convolution
tensor[d, d, tensor.size(2) // 2] = 1
elif dimensions == 4: # Spatial convolution
tensor[d, d, tensor.size(2) // 2, tensor.size(3) // 2] = 1
else: # Volumetric convolution
tensor[d, d, tensor.size(2) // 2, tensor.size(3) // 2, tensor.size(4) // 2] = 1
return tensor
with no_grad():
tensor.zero_()
for d in range(min_dim):
if dimensions == 3: # Temporal convolution
tensor[d, d, tensor.size(2) // 2] = 1
elif dimensions == 4: # Spatial convolution
tensor[d, d, tensor.size(2) // 2, tensor.size(3) // 2] = 1
else: # Volumetric convolution
tensor[d, d, tensor.size(2) // 2, tensor.size(3) // 2, tensor.size(4) // 2] = 1
return tensor
def _calculate_fan_in_and_fan_out(tensor):
......@@ -86,7 +90,9 @@ def _calculate_fan_in_and_fan_out(tensor):
num_output_fmaps = tensor.size(0)
receptive_field_size = 1
if tensor.dim() > 2:
receptive_field_size = tensor[0][0].numel()
# Fix the bug from original fb.torch
# tensor[0][0].numel() -> tensor.data[0, 0].numel()
receptive_field_size = tensor.data[0, 0].numel()
fan_in = num_input_fmaps * receptive_field_size
fan_out = num_output_fmaps * receptive_field_size
......@@ -97,13 +103,15 @@ def xavier_uniform_(tensor, gain=1):
fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
std = gain * math.sqrt(2.0 / (fan_in + fan_out))
a = math.sqrt(3.0) * std # Calculate uniform bounds from standard deviation
return tensor.uniform_(-a, a)
with no_grad():
return tensor.uniform_(-a, a)
def xavier_normal_(tensor, gain=1):
fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
std = gain * math.sqrt(2.0 / (fan_in + fan_out))
return tensor.normal_(0, std)
with no_grad():
return tensor.normal_(0, std)
def _calculate_correct_fan(tensor, mode):
......@@ -121,14 +129,16 @@ def kaiming_uniform_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu'):
gain = calculate_gain(nonlinearity, a)
std = gain / math.sqrt(fan)
bound = math.sqrt(3.0) * std # Calculate uniform bounds from standard deviation
return tensor.uniform_(-bound, bound)
with no_grad():
return tensor.uniform_(-bound, bound)
def kaiming_normal_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu'):
fan = _calculate_correct_fan(tensor, mode)
gain = calculate_gain(nonlinearity, a)
std = gain / math.sqrt(fan)
return tensor.normal_(0, std)
with no_grad():
return tensor.normal_(0, std)
# for backward compatibility
......
Dragon/python/dragon/vm/torch/nn/modules/activation.py
......@@ -32,4 +32,27 @@ class ReLU(Module):
def forward(self, x):
inputs = [x]; self.unify_devices(inputs)
outputs = [x if self._inplace else self.register_output(x.dtype)]
return self.run(inputs, outputs)
class Softmax(Module):
def __init__(self, dim=None):
super(Softmax, self).__init__()
self.dim = dim
if dim is None:
raise ValueError('Excepted a valid dim, got None.')
self.register_op()
def register_op(self):
self.op_meta = {
'op_type': 'Softmax',
'n_inputs': 1, 'n_outputs': 1,
'arguments': {
'axis': self.dim,
}
}
def forward(self, x):
inputs = [x]; self.unify_devices(inputs)
outputs = [self.register_output(x.dtype)]
return self.run(inputs, outputs)
\ No newline at end of file
Dragon/python/dragon/vm/torch/nn/modules/batchnorm.py
......@@ -20,7 +20,7 @@ from dragon.vm.torch.module import RunOperator
class _BatchNorm(Module):
def __init__(self, num_features, eps=1e-3, momentum=0.1, affine=True,
def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True,
track_running_stats=True):
super(_BatchNorm, self).__init__()
self.num_features = num_features
......
Dragon/python/dragon/vm/torch/nn/modules/container.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>
#
# Codes are based on:
#
# <https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/container.py>
#
# ------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import warnings
from collections import OrderedDict, Iterable
from itertools import islice
import operator
from dragon.vm.torch.environ import get_module_name
from dragon.vm.torch.nn import Module
class Container(Module):
def __init__(self, **kwargs):
super(Container, self).__init__()
# DeprecationWarning is ignored by default <sigh>
warnings.warn("nn.Container is deprecated. All of it's functionality "
"is now implemented in nn.Module. Subclass that instead.")
for key, value in kwargs.items():
self.add_module(key, value)
class Sequential(Module):
r"""A sequential container.
Modules will be added to it in the order they are passed in the constructor.
Alternatively, an ordered dict of modules can also be passed in.
To make it easier to understand, here is a small example::
# Example of using Sequential
model = nn.Sequential(
nn.Conv2d(1,20,5),
nn.ReLU(),
nn.Conv2d(20,64,5),
nn.ReLU()
)
# Example of using Sequential with OrderedDict
model = nn.Sequential(OrderedDict([
('conv1', nn.Conv2d(1,20,5)),
('relu1', nn.ReLU()),
('conv2', nn.Conv2d(20,64,5)),
('relu2', nn.ReLU())
]))
"""
def __init__(self, *args):
super(Sequential, self).__init__()
if len(args) == 1 and isinstance(args[0], OrderedDict):
for key, module in args[0].items():
self.add_module(key, module)
else:
for idx, module in enumerate(args):
self.add_module(str(idx), module)
def _get_item_by_idx(self, iterator, idx):
"""Get the idx-th item of the iterator"""
size = len(self)
idx = operator.index(idx)
if not -size <= idx < size:
raise IndexError('index {} is out of range'.format(idx))
idx %= size
return next(islice(iterator, idx, None))
def __getitem__(self, idx):
if isinstance(idx, slice):
return Sequential(OrderedDict(list(self._modules.items())[idx]))
else:
return self._get_item_by_idx(self._modules.values(), idx)
def __setitem__(self, idx, module):
key = self._get_item_by_idx(self._modules.keys(), idx)
return setattr(self, key, module)
def __delitem__(self, idx):
if isinstance(idx, slice):
for key in list(self._modules.keys())[idx]:
delattr(self, key)
else:
key = self._get_item_by_idx(self._modules.keys(), idx)
delattr(self, key)
def __len__(self):
return len(self._modules)
def __dir__(self):
keys = super(Sequential, self).__dir__()
keys = [key for key in keys if not key.isdigit()]
return keys
def forward(self, input):
for module in self._modules.values():
input = module(input)
return input
class ModuleList(Module):
def __init__(self, modules=None):
super(ModuleList, self).__init__()
if modules is not None:
self += modules
def _get_abs_string_index(self, idx):
"""Get the absolute index for the list of modules"""
idx = operator.index(idx)
if not (-len(self) <= idx < len(self)):
raise IndexError('index {} is out of range'.format(idx))
if idx < 0:
idx += len(self)
return str(idx)
def __getitem__(self, idx):
if isinstance(idx, slice):
return ModuleList(list(self._modules.values())[idx])
else:
return self._modules[self._get_abs_string_index(idx)]
def __setitem__(self, idx, module):
idx = operator.index(idx)
return setattr(self, str(idx), module)
def __delitem__(self, idx):
if isinstance(idx, slice):
for k in range(len(self._modules))[idx]:
delattr(self, str(k))
else:
delattr(self, self._get_abs_string_index(idx))
# To preserve numbering, self._modules is being reconstructed with modules after deletion
str_indices = [str(i) for i in range(len(self._modules))]
self._modules = OrderedDict(list(zip(str_indices, self._modules.values())))
def __len__(self):
return len(self._modules)
def __iter__(self):
return iter(self._modules.values())
def __iadd__(self, modules):
return self.extend(modules)
def __dir__(self):
keys = super(ModuleList, self).__dir__()
keys = [key for key in keys if not key.isdigit()]
return keys
def append(self, module):
# Since modules are not called from the parent module,
# We can not obtain the correct name scope
# Extend with the parent scope upon registering may be a solution
idx = str(len(self))
scope = get_module_name(self) + '/'
self.add_module(idx, module, name_v2=scope + idx)
return self
def extend(self, modules):
if not isinstance(modules, Iterable):
raise TypeError("ModuleList.extend should be called with an "
"iterable, but got " + type(modules).__name__)
offset = len(self)
scope = get_module_name(self) + '/'
for i, module in enumerate(modules):
idx = str(offset + i)
self.add_module(idx, module, name_v2=scope + idx)
return self
Dragon/python/dragon/vm/torch/nn/modules/conv.py
......@@ -102,7 +102,7 @@ class Conv2d(_ConvNd):
False, _pair(0), groups, bias)
def forward(self, input):
inputs = [input, self.weight] + [self.bias] if self.bias else []
inputs = [input, self.weight] + ([self.bias] if self.bias else [])
self.unify_devices(inputs)
outputs = [self.register_output(input.dtype)]
return self.run(inputs, outputs)
\ No newline at end of file
Dragon/python/dragon/vm/torch/nn/modules/dropout.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>
#
# ------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from dragon.vm.torch.nn import Module
class Dropout(Module):
def __init__(self, p=0.5, inplace=False):
super(Dropout, self).__init__()
self.p = p
self.inplace = inplace
self.register_op()
def register_op(self):
self.op_meta = {
'op_type': 'Dropout',
'n_inputs': 1, 'n_outputs': 1,
'arguments': {
'prob': self.p,
}
}
def forward(self, input):
if not input.requires_grad: return input
inputs = [input]
self.unify_devices(inputs)
outputs = [input if self.inplace else self.register_output(input.dtype)]
return self.run(inputs, outputs)
class Dropout2d(Dropout):
"""Dragon does not use separate backend functions."""
pass
class Dropout3d(Dropout):
"""Dragon does not use separate backend functions."""
pass
\ No newline at end of file
Dragon/python/dragon/vm/torch/nn/modules/linear.py
......@@ -49,7 +49,7 @@ class Linear(Module):
self.bias.data.uniform_(-stdv, stdv)
def forward(self, input):
inputs = [input, self.weight] + [self.bias] if self.bias else []
inputs = [input, self.weight] + ([self.bias] if self.bias else [])
self.unify_devices(inputs)
outputs = [self.register_output(input.dtype)]
return self.run(inputs, outputs)
\ No newline at end of file
Dragon/python/dragon/vm/torch/ops/__init__.py
......@@ -19,5 +19,5 @@ from .arithmetic import (
)
from .ndarray import (
sum, mean, argmin, argmax,
sum, mean, argmin, argmax, max, topk, cat
)
\ No newline at end of file
Dragon/python/dragon/vm/torch/ops/arithmetic.py
......@@ -13,39 +13,37 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
import dragon as dg
from dragon.vm.torch.tensor import Tensor
from dragon.vm.torch.tensor_uitls import from_dragon
from dragon.vm.torch.ops.primitive import MakeContext
from dragon.vm.torch.ops.primitive import MakeContext, WrapScalar
from dragon.vm.torch.ops.factory import get_module
from dragon.vm.torch.ops.modules.arithmetic import Fundamental
def _wrap_scalar(scalar, dtype):
# TODO(PhyscalX): We use (dtype/value) to hash different scalars.
# TODO(PhyscalX): In Dragon, set a Tensor with same dtype and shape will not deconstruct it.
value = np.array([scalar], dtype=dtype)
t = dg.Tensor('/share/scalar/{}/{}'.format(
dtype, str(float(scalar)))).Variable()
t.set_value(value)
return t.name
def _fundamental(input, value, op='Add', out=None):
if not isinstance(value, Tensor):
if not isinstance(value, (int, float)):
raise TypeError('Type of value should be numerical, got {}.'
.format(type(value)))
value = WrapScalar(value, input._dtype, input._ctx)
ctx = MakeContext(inputs=[input, value])
key = 'torch/ops/{}/{}:{}'.format(op.lower(), ctx[0].lower(), ctx[1])
module = get_module(Fundamental, key, ctx, op_type=op)
return module.forward(input, value, out)
def _fundamental(input, value, op='Add', out=None):
def _rfundamental(input, value, op='RAdd', out=None):
if not isinstance(value, Tensor):
if not isinstance(value, (int, float)):
raise TypeError('Type of value should be numerical, got {}.'
.format(type(value)))
# TODO(PhyscalX): We simply use a global tensor to share all scalars
scalar_name = _wrap_scalar(value, dtype=input._dtype)
value = Tensor(dg_tensor=scalar_name, ctx=input._ctx, own_storage=False)
.format(type(value)))
value = WrapScalar(value, input._dtype, input._ctx)
ctx = MakeContext(inputs=[input, value])
key = 'torch/ops/fundamental/{}:{}'.format(ctx[0].lower(), ctx[1])
key = 'torch/ops/{}/{}:{}'.format(op.lower(), ctx[0].lower(), ctx[1])
module = get_module(Fundamental, key, ctx, op_type=op)
return module(input, value, out)
return module.forward(value, input, out)
def add(input, value, out=None):
......
Dragon/python/dragon/vm/torch/ops/builtin.py
......@@ -16,10 +16,14 @@ from __future__ import print_function
from dragon.vm.torch.tensor import Tensor
from dragon.vm.torch.execute_engine import RunOperator
from dragon.vm.torch.ops.factory import get_module
from dragon.vm.torch.autograd.grad_mode import no_grad
from dragon.vm.torch.ops.primitive import MakeContext
from dragon.vm.torch.ops.arithmetic import _fundamental
from dragon.vm.torch.ops.arithmetic import _fundamental, _rfundamental
from dragon.vm.torch.ops.control_flow import _copy
from dragon.vm.torch.ops.ndarray import reshape, _fill, _reduce, _crop
from dragon.vm.torch.ops.ndarray import \
(reshape, _fill, _reduce, _arg_reduce, _crop)
from dragon.vm.torch.ops.modules.dtype import AsType
##############################################
......@@ -43,7 +47,7 @@ def copy_(self, src):
The ``self`` tensor.
"""
return _copy(src, out=self)
return _copy(self, src)
Tensor.copy_ = copy_
......@@ -69,7 +73,7 @@ def fill_(self, value):
The self.
"""
return _fill(self, self.shape, value, self)
return _fill(self, self.shape, value)
def uniform_(self, low=0, high=1):
......@@ -164,6 +168,10 @@ def add_(self, value):
return _fundamental(self, value, out=self, op='Add')
def radd(self, value):
return _rfundamental(self, value, op='RAdd')
def sub(self, value):
"""Subtract the ``self`` and ``value`` into the output tensor.
......@@ -198,6 +206,10 @@ def sub_(self, value):
return _fundamental(self, value, out=self, op='Sub')
def rsub(self, value):
return _rfundamental(self, value, op='RSub')
def mul(self, value):
"""Multiply the ``self`` and ``value`` into the output tensor.
......@@ -232,6 +244,10 @@ def mul_(self, value):
return _fundamental(self, value, out=self, op='Mul')
def rmul(self, value):
return _rfundamental(self, value, op='RMul')
def div(self, value):
"""Divide the ``self`` and ``value`` into the output tensor.
......@@ -266,14 +282,23 @@ def div_(self, value):
return _fundamental(self, value, out=self, op='Div')
def rdiv(self, value):
return _rfundamental(self, value, op='RDiv')
Tensor.add = add
Tensor.add_ = add_
Tensor.__radd__ = radd
Tensor.sub = sub
Tensor.sub_ = sub_
Tensor.__rsub__ = rsub
Tensor.mul = mul
Tensor.mul_ = mul_
Tensor.__rmul__ = rmul
Tensor.div = div
Tensor.div_ = div_
Tensor.__rdiv__ = rdiv
Tensor.__rtruediv__ = rdiv
##############################################
......@@ -302,15 +327,145 @@ def crop(self, starts, ends):
def mean(self, dim=None, keepdim=False):
return _reduce(self, 'Reduce', 'MEAN', dim, keepdim)
return _reduce(self, 'MEAN', dim, keepdim)
def sum(self, dim=None, keepdim=False):
return _reduce(self, 'Reduce', 'SUM', dim, keepdim)
return _reduce(self, 'SUM', dim, keepdim)
def max(self, dim=None, keepdim=False):
return _arg_reduce(self, 'MAX', dim, keepdim)
def min(self, dim=None, keepdim=False):
return _arg_reduce(self, 'MIN', dim, keepdim)
Tensor.view = view
Tensor.view_as = view_as
Tensor.mean = mean
Tensor.sum = sum
Tensor._crop = crop
\ No newline at end of file
Tensor.max = max
Tensor.min = min
Tensor._crop = crop
##############################################
# #
# TYPE #
# #
##############################################
def _type_to(input, dtype='float32'):
if dtype == input._dtype: return input
ctx = MakeContext(inputs=[input])
key = 'torch/ops/astype/{}:{}/dtype:{}'.format(
ctx[0].lower(), ctx[1], dtype)
module = get_module(AsType, key, ctx, dtype=dtype)
with no_grad():
return module.forward(input)
def _type(self, dtype=None):
"""Return the data type of this tensor.
If ``dtype`` is not ``None``, cast ``self`` to the new tensor.
Parameters
----------
dtype : str
The specified type.
Returns
-------
str or vm.torch.Tensor
The data type or the new tensor.
"""
if dtype is None:
return 'torch.' + self._type2str()
else:
return _type_to(self, dtype=dtype)
def _half(self):
"""Return a ``float16`` tensor with elements of ``self``.
Returns
-------
vm.torch.Tensor
The half tensor.
"""
return _type_to(self, dtype='float16')
def _float(self):
"""Return a ``float32`` tensor with elements of ``self``.
Returns
-------
vm.torch.Tensor
The float tensor.
"""
return _type_to(self, dtype='float32')
def _double(self):
"""Return a ``float64`` tensor with elements of ``self``.
Returns
-------
vm.torch.Tensor
The double tensor.
"""
return _type_to(self, dtype='float64')
def _int(self):
"""Return a ``int32`` tensor with elements of ``self``.
Returns
-------
vm.torch.Tensor
The int tensor.
"""
return _type_to(self, dtype='int32')
def _long(self):
"""Return a ``int64`` tensor with elements of ``self``.
Returns
-------
vm.torch.Tensor
The long tensor.
"""
return _type_to(self, dtype='int64')
def _byte(self):
"""Return a ``uint8`` tensor with elements of ``self``.
Returns
-------
vm.torch.Tensor
The byte tensor.
"""
return _type_to(self, dtype='uint8')
Tensor.type = _type
Tensor.half = _half
Tensor.float = _float
Tensor.double = _double
Tensor.int = _int
Tensor.long = _long
Tensor.byte = _byte
\ No newline at end of file
Dragon/python/dragon/vm/torch/ops/control_flow.py
......@@ -14,8 +14,9 @@ from dragon.vm.torch.ops.factory import get_module
from dragon.vm.torch.ops.modules.control_flow import Copy
def _copy(input, out=None):
ctx = MakeContext(inputs=[input])
key = 'torch/ops/copy/{}:{}'.format(ctx[0].lower(), ctx[1], )
def _copy(dst, src):
if id(dst) == id(src): return dst
ctx = MakeContext(inputs=[dst])
key = 'torch/ops/copy/{}:{}'.format(ctx[0].lower(), ctx[1])
module = get_module(Copy, key, ctx)
return module.forward(input, out)
\ No newline at end of file
return module.forward(dst, src)
\ No newline at end of file
Dragon/python/dragon/vm/torch/ops/modules/concat.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>
#
# ------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from dragon.vm.torch.ops.modules.base import BaseModule
class Concat(BaseModule):
def __init__(self, key, ctx, **kwargs):
super(Concat, self).__init__(key, ctx, **kwargs)
self.axis = kwargs.get('axis', 0)
self.register_arguments()
self.register_op()
def register_arguments(self):
"""No Arguments for concat op."""
pass
def register_op(self):
self.op_meta = {
'op_type': 'Concat',
'n_inputs': 1, # Ignore
'n_outputs': 1,
'arguments': {
'axis': self.axis,
}
}
def forward(self, xs, y):
inputs = xs; self.unify_devices(inputs)
outputs = [y] if y else [self.register_output(xs[0].dtype)]
return self.run(inputs, outputs)
\ No newline at end of file
Dragon/python/dragon/vm/torch/ops/modules/control_flow.py
......@@ -33,7 +33,6 @@ class Copy(BaseModule):
'arguments': {}
}
def forward(self, x, y):
inputs = [x]; self.unify_devices(inputs)
outputs = [y] if y else [self.register_output(x.dtype)]
return self.run(inputs, outputs)
\ No newline at end of file
def forward(self, dst, src):
outputs = [dst]; self.unify_devices(outputs)
return self.run([src], outputs)
\ No newline at end of file
Dragon/python/dragon/vm/torch/ops/modules/dtype.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>
#
# ------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from dragon.vm.torch.ops.modules.base import BaseModule
class AsType(BaseModule):
def __init__(self, key, ctx, **kwargs):
super(AsType, self).__init__(key, ctx, **kwargs)
self.dtype = kwargs.get('dtype', 'float32')
self.register_arguments()
self.register_op()
def register_arguments(self):
"""No Arguments for dtype op."""
pass
def register_op(self):
self.op_meta = {
'op_type': 'AsType',
'n_inputs': 1, 'n_outputs': 1,
'arguments': {
'dtype': self.dtype,
}
}
def forward(self, x):
if x.requires_grad and not x._static_shape:
raise RuntimeError("Can't change the dtype of a non-leaf tensor if requiring grad.")
inputs = [x]; self.unify_devices(inputs)
outputs = [self.register_output(self.dtype)]
y = self.run(inputs, outputs)
y.requires_grad = x.requires_grad
y._static_shape = x._static_shape
return y
\ No newline at end of file
Dragon/python/dragon/vm/torch/ops/modules/reduce.py
......@@ -19,8 +19,7 @@ from dragon.vm.torch.ops.modules.base import BaseModule
class Reduce(BaseModule):
def __init__(self, key, ctx, **kwargs):
super(Reduce, self).__init__(key, ctx, **kwargs)
self.op_type = kwargs.get('op_type', 'Reduce')
self.tag = kwargs.get('tag', None)
self.operation = kwargs.get('operation', 'SUM')
self.axis = kwargs.get('axis', -1)
self.keep_dims = kwargs.get('keep_dims', True)
self.register_arguments()
......@@ -37,10 +36,10 @@ class Reduce(BaseModule):
def register_op(self):
self.op_meta = {
'op_type': self.op_type,
'op_type': 'Reduce',
'n_inputs': 1, 'n_outputs': 1,
'arguments': {
'operation': self.tag,
'operation': self.operation,
'axis': self.axis,
'keep_dims': self.keep_dims
}
......@@ -49,4 +48,57 @@ class Reduce(BaseModule):
def forward(self, x, y):
inputs = [x]; self.unify_devices(inputs)
outputs = [y] if y else [self.register_output(x.dtype)]
return self.run(inputs, outputs)
\ No newline at end of file
return self.run(inputs, outputs)
class ArgReduce(BaseModule):
def __init__(self, key, ctx, **kwargs):
super(ArgReduce, self).__init__(key, ctx, **kwargs)
self.operation = kwargs.get('operation', 'ARGMAX')
self.axis = kwargs.get('axis', -1)
self.keep_dims = kwargs.get('keep_dims', True)
self.top_k = kwargs.get('top_k', 1)
self.register_arguments()
self.register_op()
def register_arguments(self):
"""No Arguments for reduce op.
Mutable ``axis`` and ``keep_dims`` is non-trivial for backend,
we simply hash them in the persistent key.
"""
pass
def register_op(self):
self.op_meta = {
'op_type': 'ArgReduce',
'n_inputs': 1, 'n_outputs': 1,
'arguments': {
'operation': self.operation if 'ARG' in self.operation \
else 'ARG' + self.operation,
'axis': self.axis,
'keep_dims': self.keep_dims,
'top_k': self.top_k,
}
}
def forward(self, x, y):
inputs = [x]; self.unify_devices(inputs)
if 'ARG' in self.operation:
# Return indices only
outputs = [y] if y else [self.register_output(dtype='int64')]
return self.run(inputs, outputs)
else:
if y:
if not isinstance(y, (tuple, list)):
raise TypeError('Excepted outputs as a tuple or list, got {}.'.format(type(y)))
if len(y) != 2:
raise ValueError('Excepted 2 outputs, got {}.'.format(len(y)))
outputs = [y[1], y[0]]
else: outputs = [self.register_output('int64'), self.register_output(x.dtype)]
returns = self.run(inputs, outputs)
# Return values only
if self.axis == -1: return returns[1]
# Return values and indices
return returns[1], returns[0]
\ No newline at end of file
Dragon/python/dragon/vm/torch/ops/modules/shape.py
......@@ -31,20 +31,19 @@ class Fill(BaseModule):
def register_op(self):
self.op_meta = {
'op_type': 'Fill',
'n_inputs': 1, 'n_outputs': 1,
'n_inputs': 0, 'n_outputs': 1,
'arguments': {
'value': float(self.value),
'dims_desc': [d for d in self.shape] if len(self.shape) > 0 else None,
}
}
def forward(self, x, shape, y=None):
inputs = [x]; self.unify_devices(inputs)
outputs = [y if y else self.register_output(x.dtype)]
def forward(self, x, shape):
outputs = [x]; self.unify_devices(outputs)
if shape is not None:
for ix, d in enumerate(shape):
self.set_argument_i(self.shape[ix], d)
return self.run(inputs, outputs)
return self.run([], outputs)
class Reshape(BaseModule):
......
Dragon/python/dragon/vm/torch/ops/modules/update.py
......@@ -13,6 +13,7 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import dragon.core.mpi as mpi
from dragon.vm.torch.ops.modules.base import BaseModule
......@@ -42,4 +43,40 @@ class Update(BaseModule):
}
def forward(self, param, grad):
return self.run([grad], [param], auto_grad=False)
\ No newline at end of file
self.unify_devices([param, grad])
return self.run([grad], [param], auto_grad=False)
class Collective(BaseModule):
def __init__(self, key, ctx, **kwargs):
super(Collective, self).__init__(key, ctx, **kwargs)
self.mode = kwargs.get('mode', None)
if self.mode is None:
raise ValueError('Got invalid collective mode: {}'.format(self.mode))
self.register_arguments()
self.register_op()
def register_arguments(self):
"""No arguments for update ops."""
pass
def register_op(self):
idx, group = mpi.AllowParallel()
if idx == -1:
raise RuntimeError('The mpi node({}) dost not in '
'parallel groups. \nSet it using mpi.Parallel([..]).'.format(mpi.Rank()))
mpi_comm, mpi_group = mpi.CreateGroup(root=group[0], incl=group)
self.op_meta = {
'op_type': 'CollectiveUpdate',
'n_inputs': 1, 'n_outputs': 1, # Ignore
'arguments': {
'mode': self.mode,
'comm': mpi_comm,
'group': mpi_group,
'root': group[0], # Assume the 1st node of group as root
}
}
def forward(self, grads):
self.unify_devices(grads)
return self.run(grads, grads, auto_grad=False)
\ No newline at end of file
Dragon/python/dragon/vm/torch/ops/ndarray.py
......@@ -13,11 +13,12 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from dragon.vm.torch.ops.primitive import MakeContext
from dragon.vm.torch.ops.primitive import MakeContext, CanonicalAxis
from dragon.vm.torch.ops.factory import get_module
from dragon.vm.torch.ops.modules.shape import Reshape, Fill
from dragon.vm.torch.ops.modules.reduce import Reduce
from dragon.vm.torch.ops.modules.reduce import Reduce, ArgReduce
from dragon.vm.torch.ops.modules.crop import Crop
from dragon.vm.torch.ops.modules.concat import Concat
def reshape(input, shape, shape_like=None):
......@@ -28,22 +29,33 @@ def reshape(input, shape, shape_like=None):
return module.forward(input, shape)
def _fill(input, shape, value, out=None):
def _fill(input, shape, value):
ctx = MakeContext(inputs=[input]); len_shape = len(shape)
key = 'torch/ops/fill/{}:{}/ndims:#{}/value:{}'.format(
ctx[0].lower(), ctx[1], len_shape, value)
module = get_module(Fill, key, ctx, len_shape=len_shape, value=value)
return module.forward(input, shape, out)
return module.forward(input, shape)
def _reduce(input, op, tag=None, dim=None, keepdim=False, out=None):
def _reduce(input, operation, dim=None, keepdim=False, out=None):
ctx = MakeContext(inputs=[input])
if dim is None: dim = -1
key = 'torch/ops/{}/{}:{}/dim[{}]/keep_dims:{}'.format(
op.lower() + ':{}'.format(tag.lower()) if tag else '',
if dim is None: dim = -1; keepdim = False
elif dim < 0: dim = CanonicalAxis(input, dim)
key = 'torch/ops/{}/{}:{}/dim[{}]/keep_dims:{}'.format(operation.lower(),
ctx[0].lower(), ctx[1], dim, int(keepdim))
module = get_module(Reduce, key, ctx, op_type=op,
tag=tag, axis=dim, keep_dims=keepdim)
module = get_module(Reduce, key, ctx,
operation=operation, axis=dim, keep_dims=keepdim)
return module.forward(input, out)
def _arg_reduce(input, operation, dim=None, keepdim=False, top_k=1, out=None):
ctx = MakeContext(inputs=[input])
if dim is None: dim = -1; keepdim = False
elif dim < 0: dim = CanonicalAxis(input, dim)
key = 'torch/ops/{}/{}:{}/dim[{}]/keep_dims:{}/top_k:{}'.format(operation.lower(),
ctx[0].lower(), ctx[1], dim, int(keepdim), top_k)
module = get_module(ArgReduce, key, ctx, operation=operation,
axis=dim, keep_dims=keepdim, top_k=top_k)
return module.forward(input, out)
......@@ -67,7 +79,7 @@ def mean(input, dim=None, keepdim=False, out=None):
The mean-reduced tensor.
"""
return _reduce(input, 'Reduce', 'MEAN', dim, keepdim, out)
return _reduce(input, 'MEAN', dim, keepdim, out)
def sum(input, dim=None, keepdim=False, out=None):
......@@ -87,10 +99,10 @@ def sum(input, dim=None, keepdim=False, out=None):
Returns
-------
vm.torch.Tensor
The mean-reduced tensor.
The sum-reduced tensor.
"""
return _reduce(input, 'Reduce', 'SUM', dim, keepdim, out)
return _reduce(input, 'SUM', dim, keepdim, out)
def argmax(input, dim=None, keepdim=False, out=None):
......@@ -110,14 +122,60 @@ def argmax(input, dim=None, keepdim=False, out=None):
Returns
-------
vm.torch.Tensor
The max indices.
The maximum indices.
"""
return _reduce(input, 'Argmax', None, dim, keepdim, out)
return _arg_reduce(input, 'ARGMAX', dim, keepdim, 1, out)
def max(input, dim=None, keepdim=False, out=None):
"""Return the values and indices of maximum elements along the given axis.
Parameters
----------
input : vm.torch.Tensor
The input tensor.
dim : int or None
The axis of tensor to compute sum value.
keepdim : boolean
Whether the output tensor has dim retained or not.
out : vm.torch.Tensor or None
The optional output tensor.
Returns
-------
tuple
The maximum values and indices.
"""
return _arg_reduce(input, 'MAX', dim, keepdim, 1, out)
def argmin(input, dim=None, keepdim=False, out=None):
"""Return the indices of maximum elements along the given axis.
"""Return the indices of minimum elements along the given axis.
Parameters
----------
input : vm.torch.Tensor
The input tensor.
dim : int or None
The axis of tensor to compute sum value.
keepdim : boolean
Whether the output tensor has dim retained or not.
out : vm.torch.Tensor or None
The optional output tensor.
Returns
-------
vm.torch.Tensor
The minimum indices.
"""
return _arg_reduce(input, 'ARGMIN', dim, keepdim, 1, out)
def min(input, dim=None, keepdim=False, out=None):
"""Return the values and indices of maximum elements along the given axis.
Parameters
----------
......@@ -132,11 +190,69 @@ def argmin(input, dim=None, keepdim=False, out=None):
Returns
-------
tuple
The minimum values and indices.
"""
return _arg_reduce(input, 'MIN', dim, keepdim, 1, out)
def topk(input, k, dim=None, largest=True, sorted=True, out=None):
"""Return the k largest/smallest values and indices along the given axis.
If ``dim`` is not given, the last dimension of the input is chosen.
If ``largest`` is False then the k smallest elements are returned.
Parameters
----------
input : vm.torch.Tensor
The input tensor.
k : int
The top k.
dim : int or None
The axis of tensor to compute sum value.
largest : boolean
Whether to return largest or smallest elements.
sorted : boolean
Whether to return in the sorted order.
out : vm.torch.Tensor or None
The optional output tensor.
Returns
-------
tuple
The values and indices.
"""
operation = 'MAX' if largest else 'MIN'
if dim is None: dim = input.ndimension() - 1
return _arg_reduce(input, operation, dim, True, k, out)
def cat(seq, dim=0, out=None):
"""Concatenate the inputs along the given axis.
Parameters
----------
seq : tuple or list of vm.torch.Tensor
The sequence.
dim : int
The dim to concatenate.
out : vm.torch.Tensor or None
The optional output tensor.
Returns
-------
vm.torch.Tensor
The max indices.
The output tensor.
"""
return _reduce(input, 'Argmin', None, dim, keepdim, out)
ctx = MakeContext(inputs=seq, outputs=[out] if out else [])
key = 'torch/ops/cat/{}:{}/dim:{}'.format(
ctx[0].lower(), ctx[1], dim)
module = get_module(Concat, key, ctx, axis=dim)
return module.forward(seq, out)
def _crop(input, starts, ends):
......
Dragon/python/dragon/vm/torch/ops/primitive.py
......@@ -13,6 +13,8 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
import dragon as dg
from dragon.vm.torch.tensor import *
......@@ -25,10 +27,10 @@ def CheckDataType(inputs, dtypes=None):
input._dtype not in dtypes:
raise TypeError('Type of input({}) is {}, '
'not in the support set: ({}).'
.format(ix, input._dtype, ', '.join(dtypes)))
.format(ix, input._dtype, ', '.join(dtypes)))
if input._dtype != request_type:
raise TypeError('Excepted the type of input({}) is {}, got {}.'
.format(ix, request_type, input._dtype))
.format(ix, request_type, input._dtype))
def UnifyDevices(tensors, key='Inputs'):
......@@ -36,12 +38,12 @@ def UnifyDevices(tensors, key='Inputs'):
devices = [0]
if len(set(device_types)) != 1:
raise ValueError('{} from different device type: [{}].'
.format(key, ', '.join(device_types)))
.format(key, ', '.join(device_types)))
if device_types[0] == 'CUDA':
devices = [t._ctx[1] for t in tensors]
if len(set(devices)) != 1:
raise ValueError('{} from different cuda device: [{}].'
.format(key, ', '.join([str(d) for d in devices])))
.format(key, ', '.join([str(d) for d in devices])))
return device_types[0], devices[0]
......@@ -57,4 +59,24 @@ def MakeContext(inputs=(), outputs=(), meta=None):
# Case #3: [], [...] -> Refer Outputs
# Case #4: [...], [...] -> Refer Outputs
# Case #5: meta -> CPU, CUDA:?
return type, device_id
\ No newline at end of file
return type, device_id
def WrapScalar(scalar, dtype, ctx):
# We use (DType + Value) to hash different scalars
# Setting a Tensor with same DType and shape will not deconstruct it
value = np.array([scalar], dtype=dtype)
if 'float' in dtype: scalar = float(scalar)
if 'int' in dtype: scalar = int(scalar)
t = dg.Tensor('/share/scalar/{}/{}'.format(
dtype, str(scalar))).Variable()
t.set_value(value)
t = Tensor(dg_tensor=t.name, dtype=dtype, ctx=ctx, own_storage=False)
t.requires_grad = False
return t
def CanonicalAxis(input, dim):
ndim = input.ndimension()
while dim < 0: dim += ndim
return dim
Dragon/python/dragon/vm/torch/ops/update.py
......@@ -15,26 +15,20 @@ from __future__ import print_function
import dragon.core.mpi as mpi
from dragon.vm.torch.execute_engine import RunOperator
from dragon.vm.torch.ops.primitive import MakeContext
from dragon.vm.torch.ops.factory import get_module
from dragon.vm.torch.ops.modules.update import Update
from dragon.vm.torch.ops.modules.update import Update, Collective
def _allreduce_op(params, grads):
def _allreduce(grads):
if not mpi.Is_Init(): return
idx, group = mpi.AllowParallel()
arguments = {}
if idx != -1:
arguments['mode'] = mpi.GetParallelMode() + '_ALLREDUCE'
arguments['comm'], arguments['group'] \
= mpi.CreateGroup(root=group[0], incl=group)
arguments['root'] = group[0]
inputs = grads if isinstance(grads, list) else [grads]
outputs = params if isinstance(params, list) else [params]
ctx = MakeContext(inputs, outputs)
return RunOperator(inputs, outputs, op_type='CollectiveUpdate',
ctx=ctx, auto_grad=False, **arguments)
if not isinstance(grads, (list, tuple)): grads = [grads]
ctx = MakeContext(inputs=grads)
mode = mpi.GetParallelMode() + '_ALLREDUCE'
key = 'torch/ops/collective/{}:{}/{}'.format(
ctx[0].lower(), ctx[1], mode.lower())
module = get_module(Collective, key, ctx, mode=mode)
return module.forward(grads)
def _update(param, grad, op_type, slot,
......
Dragon/python/dragon/vm/torch/optim/optimizer.py
......@@ -24,7 +24,7 @@ import numpy as np
import dragon as dg
from dragon.vm.torch.tensor import Tensor
from dragon.vm.torch.ops.update import _allreduce_op, _update
from dragon.vm.torch.ops.update import _allreduce, _update
_OPTIMIZER_GROUP_UID = 0
......@@ -95,14 +95,15 @@ class Optimizer(object):
for p in group['params']:
g_name = p.name + '_grad'
if not dg.workspace.HasTensor(g_name): continue
g = Tensor(dg_tensor=g_name); g._own_storage = False
g = Tensor(dg_tensor=g_name)
g._own_storage = False; g._ctx = p._ctx
params.append(p); grads.append(g)
# Feed optimizer parameters to workspace
self.feed_parameters(group)
# Run a all-reduce op to accumulate grads if necessary
_allreduce_op(params, grads)
_allreduce(grads)
# Run regular update ops
for p, g in zip(params, grads):
......
Dragon/python/dragon/vm/torch/serialization.py
......@@ -79,22 +79,28 @@ def save(obj, f, pickle_module=pickle, pickle_protocol=DEFAULT_PROTOCOL):
lambda f: _save(obj, f, pickle_module, pickle_protocol))
def _load(f, map_location=None, pickle_module=pickle):
def _load(f, map_location=None, pickle_module=pickle, file=None):
try:
return pickle.load(f)
return pickle_module.load(f)
except UnicodeDecodeError:
return pickle.load(f, encoding='iso-8859-1')
if file:
# ReOpen the file, because the MARK is corrupted
f = open(file, 'rb')
return pickle_module.load(f, encoding='iso-8859-1')
else: return pickle_module.load(f, encoding='iso-8859-1')
def load(f, map_location=None, pickle_module=pickle):
new_fd = False
file = None
if isinstance(f, str) or \
(sys.version_info[0] == 2 and isinstance(f, unicode)) or \
(sys.version_info[0] == 3 and isinstance(f, pathlib.Path)):
new_fd = True
file = f
f = open(f, 'rb')
try:
return _load(f, map_location, pickle_module)
return _load(f, map_location, pickle_module, file)
finally:
if new_fd:
f.close()
\ No newline at end of file
Dragon/python/dragon/vm/torch/tensor.py
......@@ -315,6 +315,17 @@ class Tensor(object):
"""
return self.div_(other)
def __neg__(self):
"""Calculate -x.
Returns
-------
vm.torch.Tensor
The output tensor.
"""
return self.mul(-1.0)
def __repr__(self):
"""Return a format str representing the internal storage.
......@@ -404,6 +415,9 @@ class Tensor(object):
if not keep_dims: ends[-1] = -1
return self._crop(starts, ends)
def __setitem__(self, key, value):
print(key, value.name)
def device(self):
return self._ctx[1]
......@@ -489,16 +503,24 @@ class Tensor(object):
"""
return self._dtype
def type(self):
def type(self, dtype=None):
"""Return the data type of this tensor.
If ``dtype`` is not ``None``, cast ``self`` to the new tensor.
Parameters
----------
dtype : str
The specified type.
Returns
-------
str
The data type.
str or vm.torch.Tensor
The data type or the new tensor.
"""
return 'torch.' + self._type2str()
raise NotImplementedError('Refer torch.ops.builtin.type')
##############################################
# #
......@@ -520,7 +542,7 @@ class Tensor(object):
The output tensor.
"""
raise NotImplementedError('Refer torch.ops.builtins.view')
raise NotImplementedError('Refer torch.ops.builtin.view')
def view_as(self, other):
"""Returns a new tensor with the same data but a different size as the given tensor.
......@@ -536,7 +558,7 @@ class Tensor(object):
The output tensor.
"""
raise NotImplementedError('Refer torch.ops.builtins.view_as')
raise NotImplementedError('Refer torch.ops.builtin.view_as')
def copy_(self, src):
"""Copy the elements from ``src`` into this tensor and return ``self``.
......@@ -552,7 +574,7 @@ class Tensor(object):
The ``self`` tensor.
"""
raise NotImplementedError('Refer torch.ops.builtins.copy_')
raise NotImplementedError('Refer torch.ops.builtin.copy_')
def fill_(self, value):
"""Fills self tensor with the specified value.
......@@ -567,7 +589,7 @@ class Tensor(object):
The self.
"""
raise NotImplementedError('Refer torch.ops.builtins.fill_')
raise NotImplementedError('Refer torch.ops.builtin.fill_')
def zero_(self):
"""Fills self tensor with zeros.
......@@ -600,7 +622,7 @@ class Tensor(object):
The self.
"""
raise NotImplementedError('Refer torch.ops.builtins.uniform_')
raise NotImplementedError('Refer torch.ops.builtin.uniform_')
def normal_(self, mean=0, std=1):
"""Fill self tensor with the specified normal distribution.
......@@ -618,7 +640,7 @@ class Tensor(object):
The self.
"""
raise NotImplementedError('Refer torch.ops.builtins.normal_')
raise NotImplementedError('Refer torch.ops.builtin.normal_')
def add(self, value):
"""See ``torch.add()``
......@@ -634,7 +656,7 @@ class Tensor(object):
The output tensor.
"""
raise NotImplementedError('Refer torch.ops.builtins.add')
raise NotImplementedError('Refer torch.ops.builtin.add')
def add_(self, value):
"""Inplace of ``torch.add()``
......@@ -650,7 +672,7 @@ class Tensor(object):
The self.
"""
raise NotImplementedError('Refer torch.ops.builtins.add_')
raise NotImplementedError('Refer torch.ops.builtin.add_')
def sub(self, value):
"""Subtract the ``self`` and ``value`` into the output tensor.
......@@ -666,7 +688,7 @@ class Tensor(object):
The output tensor.
"""
raise NotImplementedError('Refer torch.ops.builtins.sub')
raise NotImplementedError('Refer torch.ops.builtin.sub')
def sub_(self, value):
"""Inplace of ``Tensor.sub()``
......@@ -682,7 +704,7 @@ class Tensor(object):
The self.
"""
raise NotImplementedError('Refer torch.ops.builtins.sub_')
raise NotImplementedError('Refer torch.ops.builtin.sub_')
def mul(self, value):
"""Multiply the ``self`` and ``value`` into the output tensor.
......@@ -698,7 +720,7 @@ class Tensor(object):
The output tensor.
"""
raise NotImplementedError('Refer torch.ops.builtins.mul')
raise NotImplementedError('Refer torch.ops.builtin.mul')
def mul_(self, value):
"""Inplace of ``Tensor.mul()``
......@@ -714,7 +736,7 @@ class Tensor(object):
The self.
"""
raise NotImplementedError('Refer torch.ops.builtins.mul_')
raise NotImplementedError('Refer torch.ops.builtin.mul_')
def div(self, value):
"""Divide the ``self`` and ``value`` into the output tensor.
......@@ -730,7 +752,7 @@ class Tensor(object):
The output tensor.
"""
raise NotImplementedError('Refer torch.ops.builtins.div')
raise NotImplementedError('Refer torch.ops.builtin.div')
def div_(self, value):
"""Inplace of ``Tensor.div()``
......@@ -746,7 +768,7 @@ class Tensor(object):
The self.
"""
raise NotImplementedError('Refer torch.ops.builtins.div_')
raise NotImplementedError('Refer torch.ops.builtin.div_')
def mean(self, dim=None, keepdim=False):
"""Returns the mean of all elements or elements along the given dim.
......@@ -764,9 +786,9 @@ class Tensor(object):
The mean-reduced tensor.
"""
raise NotImplementedError('Refer torch.ops.builtins.mean')
raise NotImplementedError('Refer torch.ops.builtin.mean')
def sum(input, dim=None, keepdim=False):
def sum(self, dim=None, keepdim=False):
"""Returns the sum of all elements or elements along the given dim.
Parameters
......@@ -779,10 +801,112 @@ class Tensor(object):
Returns
-------
vm.torch.Tensor
The mean-reduced tensor.
The sum-reduced tensor.
"""
raise NotImplementedError('Refer torch.ops.builtin.sum')
def max(self, dim=None, keepdim=False):
"""Return the values and indices of maximum elements along the given axis.
Parameters
----------
dim : int or None
The axis of tensor to compute sum value.
keepdim : boolean
Whether the output tensor has dim retained or not.
Returns
-------
vm.torch.Tensor
The maximum values and indices.
"""
raise NotImplementedError('Refer torch.ops.builtin.max')
def min(input, dim=None, keepdim=False):
"""Return the values and indices of minimum elements along the given axis.
Parameters
----------
dim : int or None
The axis of tensor to compute sum value.
keepdim : boolean
Whether the output tensor has dim retained or not.
Returns
-------
vm.torch.Tensor
The minimum values and indices.
"""
raise NotImplementedError('Refer torch.ops.builtin.min')
def half(self):
"""Return a ``float16`` tensor with elements of ``self``.
Returns
-------
vm.torch.Tensor
The half tensor.
"""
raise NotImplementedError('Refer torch.ops.builtin.half')
def float(self):
"""Return a ``float32`` tensor with elements of ``self``.
Returns
-------
vm.torch.Tensor
The float tensor.
"""
raise NotImplementedError('Refer torch.ops.builtin.float')
def double(self):
"""Return a ``float64`` tensor with elements of ``self``.
Returns
-------
vm.torch.Tensor
The float tensor.
"""
raise NotImplementedError('Refer torch.ops.builtins.sum')
raise NotImplementedError('Refer torch.ops.builtin.double')
def int(self):
"""Return a ``int32`` tensor with elements of ``self``.
Returns
-------
vm.torch.Tensor
The int tensor.
"""
raise NotImplementedError('Refer torch.ops.builtin.int')
def long(self):
"""Return a ``int64`` tensor with elements of ``self``.
Returns
-------
vm.torch.Tensor
The long tensor.
"""
raise NotImplementedError('Refer torch.ops.builtin.long')
def byte(self):
"""Return a ``uint8`` tensor with elements of ``self``.
Returns
-------
vm.torch.Tensor
The byte tensor.
"""
raise NotImplementedError('Refer torch.ops.builtin.byte')
##############################################
# #
......@@ -810,7 +934,7 @@ class Tensor(object):
@property
def grad(self):
g = from_dragon(self.name + '_grad', False)
g._static_shape = self.shape
if g: g._static_shape = self.shape
return g
@property
......@@ -836,11 +960,19 @@ class Tensor(object):
##############################################
def _type2str(self):
return {'float32': 'FloatTensor',
'float64': 'DoubleTensor',
'int32': 'IntTensor',
'int64': 'LongTensor',
'uint8': 'ByteTensor'}[self._dtype]
return {
'float16': 'HalfTensor',
'float32': 'FloatTensor',
'float64': 'DoubleTensor',
'int32': 'IntTensor',
'int64': 'LongTensor',
'uint8': 'ByteTensor'
}[self._dtype]
def HalfTensor(*args, **kwargs):
kwargs['dtype'] = 'float16'
return Tensor(*args, **kwargs)
def FloatTensor(*args, **kwargs):
......@@ -869,6 +1001,7 @@ def ByteTensor(*args, **kwargs):
_DTYPE_TO_TENSOR = {
'float16': HalfTensor,
'float32': FloatTensor,
'float64': DoubleTensor,
'int32': IntTensor,
......@@ -887,7 +1020,7 @@ def LeafTensor(shape, dtype='float32', ctx=None, requires_grad=False):
return constructor(*shape, ctx=ctx, requires_grad=requires_grad)
def RuntimeTensor(dtype='float32', ctx=None):
def RuntimeTensor(name, dtype='float32', ctx=None):
"""Create a torch tensor according to dtype and ctx.
Commonly used to represent the outputs that are hard to compute shape,
......@@ -895,7 +1028,6 @@ def RuntimeTensor(dtype='float32', ctx=None):
"""
constructor = _DTYPE_TO_TENSOR[dtype]
name = TPool.get('runtime')
dg.workspace.CreateTensor(name)
return constructor(dg_tensor=name, ctx=ctx)
......
Dragon/python/dragon/vm/torch/tensor_pool.py
......@@ -14,6 +14,16 @@
Tensors with the same scope in the pool will be reused by turns,
which speeds up the whole system by reducing the unnecessary deconstructing.
Heuristically, we have used 4 pools with different scopes:
* scope(leaf): A Pool to reuse leaf tensors.
* scope(numpy): A pool to reuse leaf tensors from numpy.
* scope(join): A pool to reuse RT(runtime) tensors required by forward-backward.
* scope(detach): A pool to reuse RT(runtime) tensors required by forward only.
"""
from __future__ import absolute_import
......@@ -43,7 +53,7 @@ class TensorPool(object):
self._scope2handle[scope] - 1)
return self._scope2keys[scope].popleft()
def get(self, scope='runtime'):
def get(self, scope='detach'):
handle = self.get_handle(scope)
return '[TPool]{}/tensor:{}'.format(scope, handle)
......@@ -55,7 +65,7 @@ class TensorPool(object):
scope, handle = name[7:].split('/tensor:')
if not handle.isdigit():
raise ValueError('Got illegal name: ' + name + '.\n'
'Excepted the format: [TPool]/prefix/tensor:handle')
'Excepted the format: [TPool]/scope/tensor:handle')
self.put_handle(scope, handle)
return True
else: return False
......
Dragon/python/dragon/vm/torch/tensor_uitls.py
......@@ -76,7 +76,7 @@ def from_dragon(tensor, own_storage=False):
"""
info = dg_tensor_utils.GetTensorInfo(tensor)
if not info['init']: return None
if not info or not info['init']: return None
module = importlib.import_module('dragon.vm.torch.tensor')
th_tensor = getattr(module, type_np2torch(info['dtype']))()
th_tensor._ctx = (info['mem_at'], info['device_id'])
......
Dragon/python/dragon/vm/torch/tools/__init__.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>
#
# ------------------------------------------------------------
\ No newline at end of file
Dragon/python/dragon/vm/torch/tools/model_converter.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>
#
# ------------------------------------------------------------
"""Please convert models at python2.7,
as the protocols of python3.x of pickle are not downward compatible.
"""
import os
import sys
from collections import OrderedDict
import torch
from torch.utils.model_zoo import \
_download_url_to_file, urlparse, HASH_REGEX
def state_dict_v2(m, destination=None, prefix=''):
"""A forked version of nn.Module.state_dict().
We apply a tensor2numpy for each parameters recursively,
which is compatible with dragon.vm.torch.nn.Module.load_state_dict().
Parameters
----------
m : torch.nn.Module
The fb.pytorch module.
destination : OrderedDict or None
The output dict.
prefix : str
The prefix to the key of dict.
Returns
-------
OrderedDict
The dict with numpy parameters.
"""
t2np = lambda t : t.cpu().numpy()
if destination is None:
destination = OrderedDict()
destination._metadata = OrderedDict()
for name, param in m._parameters.items():
if param is not None:
destination[prefix + name] = t2np(param.data)
for name, buf in m._buffers.items():
if buf is not None:
destination[prefix + name] = t2np(buf)
for name, module in m._modules.items():
if module is not None:
state_dict_v2(module, destination, prefix + name + '.')
return destination
if __name__ == '__main__':
from torchvision.models import (
alexnet,
vgg11, vgg13, vgg16, vgg19, vgg11_bn, vgg13_bn, vgg16_bn, vgg19_bn,
resnet18, resnet34, resnet50, resnet101, resnet152,
squeezenet1_0, squeezenet1_1,
inception_v3,
)
tasks = OrderedDict([
# (alexnet, 'https://download.pytorch.org/models/alexnet-owt-4df8aa71.pth'),
(vgg11, 'https://download.pytorch.org/models/vgg11-bbd30ac9.pth'),
(vgg13, 'https://download.pytorch.org/models/vgg13-c768596a.pth'),
# (vgg16, 'https://download.pytorch.org/models/vgg16-397923af.pth'),
# (vgg16, 'https://download.pytorch.org/models/vgg16-397923af.pth'),
# (vgg16, 'https://download.pytorch.org/models/vgg16-397923af.pth'),
# (vgg19, 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth'),
(vgg11_bn, 'https://download.pytorch.org/models/vgg11_bn-6002323d.pth'),
(vgg13_bn, 'https://download.pytorch.org/models/vgg13_bn-abd245e5.pth'),
(vgg16_bn, 'https://download.pytorch.org/models/vgg16_bn-6c64b313.pth'),
(vgg19_bn, 'https://download.pytorch.org/models/vgg19_bn-c79401a0.pth'),
# (resnet18, 'https://download.pytorch.org/models/resnet18-5c106cde.pth'),
# (resnet34, 'https://download.pytorch.org/models/resnet34-333f7ec4.pth'),
# (resnet50, 'https://download.pytorch.org/models/resnet50-19c8e357.pth'),
# (resnet101, 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth'),
# (resnet152, 'https://download.pytorch.org/models/resnet152-b121ed2d.pth'),
# (squeezenet1_0, 'https://download.pytorch.org/models/squeezenet1_0-a815701f.pth'),
# (squeezenet1_1, 'https://download.pytorch.org/models/squeezenet1_1-f364aa15.pth'),
# (inception_v3, 'https://download.pytorch.org/models/inception_v3_google-1a9a5a14.pth'),
])
downloads = []
torch_home = os.path.expanduser(os.getenv('TORCH_HOME', '~/.torch'))
model_dir = os.getenv('TORCH_MODEL_ZOO', os.path.join(torch_home, 'models'))
for m, url in tasks.items():
if not os.path.exists(model_dir): os.makedirs(model_dir)
parts = urlparse(url)
filename = os.path.basename(parts.path)
cached_file = os.path.join(model_dir, filename)
if not os.path.exists(cached_file):
sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file))
hash_prefix = HASH_REGEX.search(filename).group(1)
_download_url_to_file(url, cached_file, hash_prefix, progress=True)
downloads.append((m, cached_file))
if sys.version_info >= (3,0):
raise RuntimeError('You can download with python2|3, but convert with python2 only.')
import cPickle
for m, file in downloads:
p1, p2 = os.path.split(file)
p3, p4 = os.path.splitext(p2)
file_v2 = os.path.join(p1, p3 + '.pkl')
mi = m()
mi.load_state_dict(torch.load(file))
np_state_dict = state_dict_v2(mi)
cPickle.dump(np_state_dict, open(file_v2, 'wb'), cPickle.HIGHEST_PROTOCOL)
print('Convert {} to {}.'.format(file, file_v2))
\ No newline at end of file
Dragon/python/dragon/vm/torch/utils/__init__.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>
#
# ------------------------------------------------------------
\ No newline at end of file
Dragon/python/dragon/vm/torch/utils/data/__init__.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>
#
# ------------------------------------------------------------
from .dataset import Dataset
from .dataloader import DataLoader
\ No newline at end of file
Dragon/python/dragon/vm/torch/utils/data/dataloader.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>
#
# ------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from dragon.vm.torch.utils.data.io.data_batch import DataBatch as _DataBatch
class _DataLoaderIter(object):
def __init__(self, loader):
self.loader = loader
def __len__(self):
return len(self.loader.batch.Q_level_3.qsize())
def __next__(self):
return self.loader.batch.get()
next = __next__ # Python 2 compatibility
def __iter__(self):
return self
class DataLoader(object):
def __init__(self, dataset, batch_size=1, shuffle=False,
partition=False, multiple_nodes=False, instance_chunk=False):
"""A MPI-Aware DataLoader. Forked from ``dragon.io``.
Parameters
----------
dataset : vm.torch.utils.data.dataset.Dataset
The dataset.
batch_size : int
The batch size. Divided by n mpi-nodes if ``partition`` is True.
instance_chunk : boolean
Whether to limit each chunk with at most 1 instance.
shuffle : boolean
Whether to shuffle the data.
partition : boolean
Whether to partition batch. Default is ``False``.
multiple_nodes: boolean
Whether to split data for multiple parallel nodes.
instance_chunk : boolean
Whether to limit each chunk with at most 1 instance.
"""
self.dataset = dataset
self.batch_size = batch_size
n_transformers = 1
if dataset.transform and \
hasattr(dataset.transform, 'n_transformers'):
n_transformers = dataset.transform.n_transformers
self.batch = _DataBatch(**{
'source': dataset.database,
'multiple_nodes': multiple_nodes,
'shuffle': shuffle,
'instance_chunk': instance_chunk,
'batch_size': batch_size,
'partition': partition,
'transform': dataset.transform,
'color_space': dataset.color_space,
'num_transformers': n_transformers,
})
def __iter__(self):
return _DataLoaderIter(self)
def __len__(self):
return self.batch.Q_level_3.qsize()
def next(self):
return self.batch.get()
def get(self):
return self.batch.get()
\ No newline at end of file
Dragon/python/dragon/vm/torch/utils/data/dataset.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>
#
# ------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from dragon.vm.torch.utils.data.io.data_reader import DataReader
from dragon.vm.torch.utils.data.io.data_transformer import DataTransformer
class _LMDBStream(object):
def __init__(self, database, transform, color_space='RGB'):
from multiprocessing import Queue
self.Q = Queue(1)
# Create a DataReader
self._data_reader = DataReader(**{'source': database})
self._data_reader.Q_out = Queue(1)
self._data_reader.start()
# Create a DataTransformer
self._data_transformer = DataTransformer(transform=transform,
color_space=color_space, pack=True)
self._data_transformer.Q_in = self._data_reader.Q_out
self._data_transformer.Q_out = self.Q
self._data_transformer.start()
def cleanup():
def terminate(process):
process.terminate()
process.join()
terminate(self._data_transformer)
terminate(self._data_reader)
import atexit
atexit.register(cleanup)
def get(self):
return self.Q.get()
def next(self):
return self.get()
class Dataset(object):
"""An abstract class representing a Dataset.
Parameters
----------
database : str
The path of LMDB database.
transform : lambda
The transforms.
color_space : str
The color space.
"""
def __init__(self, database, transform=None, color_space='RGB'):
self.database = database
self.transform = transform
self.color_space = color_space
def create_lmdb_stream(self):
return _LMDBStream(self.database, self.transform, self.color_space)
\ No newline at end of file
Dragon/python/dragon/vm/torch/utils/data/io/__init__.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>
#
# ------------------------------------------------------------
\ No newline at end of file
Dragon/python/dragon/vm/torch/utils/data/io/blob_fetcher.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>
#
# ------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
from multiprocessing import Process
class BlobFetcher(Process):
"""BlobFetcher is deployed to queue blobs from `DataTransformer`_.
It is supported to form ``NCHW`` image blobs and ``1D`` label blobs.
"""
def __init__(self, **kwargs):
"""Construct a ``BlobFetcher``.
Parameters
----------
batch_size : int
The size of a training batch.
partition : boolean
Whether to partition batch. Default is ``False``.
prefetch : int
The prefetch count. Default is ``5``.
"""
super(BlobFetcher, self).__init__()
self._batch_size = kwargs.get('batch_size', 100)
self._partition = kwargs.get('partition', False)
if self._partition:
self._batch_size = int(self._batch_size / kwargs['group_size'])
self.Q_in = self.Q_out = None
self.daemon = True
def get(self):
"""Return a batch with image and label blob.
Returns
-------
tuple
The blob of image and labels.
"""
# fill blobs
im, labels = self.Q_in.get()
im_blob = np.zeros(shape=([self._batch_size] + list(im.shape)), dtype=im.dtype)
label_blob = np.zeros((self._batch_size, len(labels)), dtype=np.int64)
for ix in range(0, self._batch_size):
im_blob[ix, :, :, :], label_blob[ix, :] = im, labels
if ix != self._batch_size - 1: im, labels = self.Q_in.get()
# mean subtraction & numerical scale
im_blob = im_blob.astype(np.float32)
return im_blob, label_blob
def run(self):
"""Start the process.
Returns
-------
None
"""
while True:
self.Q_out.put(self.get())
\ No newline at end of file
Dragon/python/dragon/vm/torch/utils/data/io/data_batch.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>
#
# ------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import time
import pprint
from multiprocessing import Queue
import dragon.core.mpi as mpi
from .data_reader import DataReader
from .data_transformer import DataTransformer
from .blob_fetcher import BlobFetcher
class DataBatch(object):
"""DataBatch aims to prefetch data by ``Triple-Buffering``.
It takes full advantages of the Process/Thread of Python,
which provides remarkable I/O speed up for scalable distributed training.
"""
def __init__(self, **kwargs):
"""Construct a ``DataBatch``.
Parameters
----------
source : str
The path of database.
multiple_nodes: boolean
Whether to split data for multiple parallel nodes.
shuffle : boolean
Whether to shuffle the data.
instance_chunk : boolean
Whether to limit each chunk with at most 1 instance.
num_chunks : int
The number of chunks to split. Default is ``2048``.
chunk_size : int
The size(MB) of each chunk. Default is -1 (Refer ``num_chunks``).
mean_values : list
The mean value of each image channel.
scale : float
The scale performed after mean subtraction. Default is ``1.0``.
padding : int
The zero-padding size. Default is ``0`` (Disabled).
fill_value : int
The value to fill when padding is valid. Default is ``127``.
crop_size : int
The crop size. Default is ``0`` (Disabled).
mirror : boolean
Whether to flip(horizontally) images. Default is ``False``.
color_augmentation : boolean
Whether to distort colors. Default is ``False``.
min_random_scale : float
The min scale of the input images. Default is ``1.0``.
max_random_scale : float
The max scale of the input images. Default is ``1.0``.
force_color : boolean
Set to duplicate channels for gray. Default is ``False``.
phase : str
The phase of this operator, ``TRAIN`` or ``TEST``. Default is ``TRAIN``.
batch_size : int
The size of a training batch.
partition : boolean
Whether to partition batch. Default is ``False``.
prefetch : int
The prefetch count. Default is ``5``.
"""
super(DataBatch, self).__init__()
# init mpi
global_rank = 0; local_rank = 0; group_size = 1
if mpi.Is_Init():
idx, group = mpi.AllowParallel()
if idx != -1: # data parallel
global_rank = mpi.Rank()
group_size = len(group)
for i, node in enumerate(group):
if global_rank == node: local_rank = i
kwargs['group_size'] = group_size
# configuration
self._prefetch = kwargs.get('prefetch', 5)
self._num_readers = kwargs.get('num_readers', 1)
self._num_transformers = kwargs.get('num_transformers', -1)
self._max_transformers = kwargs.get('max_transformers', 3)
self._num_fetchers = kwargs.get('num_fetchers', 1)
# io-aware policy
if self._num_transformers == -1:
self._num_transformers = 1
# add 1 transformer for color augmentation
if kwargs.get('color_augmentation', False):
self._num_transformers += 1
# add 1 transformer for random scale
if kwargs.get('max_random_scale', 1.0) - \
kwargs.get('min_random_scale', 1.0) != 0:
self._num_transformers += 1
self._num_transformers = min(self._num_transformers, self._max_transformers)
self._batch_size = kwargs.get('batch_size', 100)
self._partition = kwargs.get('partition', False)
if self._partition:
self._batch_size = int(self._batch_size / kwargs['group_size'])
# init queues
self.Q_level_1 = Queue(self._prefetch * self._num_readers * self._batch_size)
self.Q_level_2 = Queue(self._prefetch * self._num_readers * self._batch_size)
self.Q_level_3 = Queue(self._prefetch * self._num_readers)
# init readers
self._readers = []
for i in range(self._num_readers):
self._readers.append(DataReader(**kwargs))
self._readers[-1].Q_out = self.Q_level_1
for i in range(self._num_readers):
num_parts = self._num_readers
part_idx = i
if self._readers[i]._multiple_nodes or \
self._readers[i]._use_shuffle:
num_parts *= group_size
part_idx += local_rank * self._num_readers
self._readers[i]._num_parts = num_parts
self._readers[i]._part_idx = part_idx
self._readers[i]._random_seed += part_idx
self._readers[i].start()
time.sleep(0.1)
# init transformers
self._transformers = []
for i in range(self._num_transformers):
transformer = DataTransformer(**kwargs)
transformer._random_seed += (i + local_rank * self._num_transformers)
transformer.Q_in = self.Q_level_1
transformer.Q_out = self.Q_level_2
transformer.start()
self._transformers.append(transformer)
time.sleep(0.1)
# init blob fetchers
self._fetchers = []
for i in range(self._num_fetchers):
fetcher = BlobFetcher(**kwargs)
fetcher.Q_in = self.Q_level_2
fetcher.Q_out = self.Q_level_3
fetcher.start()
self._fetchers.append(fetcher)
time.sleep(0.1)
# prevent to echo multiple nodes
if local_rank == 0: self.echo()
def cleanup():
def terminate(processes):
for process in processes:
process.terminate()
process.join()
from dragon.config import logger
terminate(self._fetchers)
if local_rank == 0: logger.info('Terminating BlobFetcher ......')
terminate(self._transformers)
if local_rank == 0: logger.info('Terminating DataTransformer ......')
terminate(self._readers)
if local_rank == 0: logger.info('Terminating DataReader......')
import atexit
atexit.register(cleanup)
def get(self):
"""Get a batch.
Returns
-------
tuple
The batch, representing data and labels respectively.
"""
return self.Q_level_3.get()
def echo(self):
"""Print I/O Information.
Returns
-------
None
"""
print('---------------------------------------------------------')
print('BatchFetcher({} Threads), Using config:'.format(
self._num_readers + self._num_transformers + self._num_fetchers))
params = {'queue_size': self._prefetch,
'n_readers': self._num_readers,
'n_transformers': self._num_transformers,
'n_fetchers': self._num_fetchers}
pprint.pprint(params)
print('---------------------------------------------------------')
Dragon/python/dragon/vm/torch/utils/data/io/data_reader.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>
#
# ------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import math
import numpy as np
import numpy.random as npr
from multiprocessing import Process
import dragon.config as config
from dragon.tools.db import LMDB
class DataReader(Process):
"""DataReader is deployed to queue encoded str from `LMDB`_.
It is supported to adaptively partition and shuffle records over all distributed nodes.
"""
def __init__(self, **kwargs):
"""Construct a ``DataReader``.
Parameters
----------
source : str
The path of database.
multiple_nodes: boolean
Whether to split data for multiple parallel nodes.
shuffle : boolean
Whether to shuffle the data.
instance_chunk : boolean
Whether to limit each chunk with at most 1 instance.
num_chunks : int
The number of chunks to split. Default is ``2048``.
chunk_size : int
The size(MB) of each chunk. Default is -1 (Refer ``num_chunks``).
"""
super(DataReader, self).__init__()
self._source = kwargs.get('source', '')
self._multiple_nodes = kwargs.get('multiple_nodes', False)
self._use_shuffle = kwargs.get('shuffle', False)
self._use_instance_chunk = kwargs.get('instance_chunk', False)
self._num_chunks = kwargs.get('num_chunks', 2048)
self._chunk_size = kwargs.get('chunk_size', -1)
self._num_parts = 1
self._part_idx = 0
self._random_seed = config.GetRandomSeed()
self._cur_idx = 0
self._cur_chunk_idx = 0
self.Q_out = None
self.daemon = True
def element(self):
"""Get the value of current record.
Returns
-------
str
The encoded str.
"""
return self._db.value()
def redirect(self, target_idx):
"""Redirect to the target position.
Parameters
----------
target_idx : int
The key of instance in ``LMDB``.
Returns
-------
None
Notes
-----
The redirection reopens the ``LMDB``.
You can drop caches by ``echo 3 > /proc/sys/vm/drop_caches``.
This will disturb getting stuck when ``Database Size`` >> ``RAM Size``.
"""
self._db.close()
self._db.open(self._source)
self._cur_idx = target_idx
self._db.set(str(self._cur_idx).zfill(self._db_zfill))
def reset(self):
"""Reset the cursor and environment.
Returns
-------
None
"""
if self._multiple_nodes or self._use_shuffle:
if self._use_shuffle: self._perm = npr.permutation(self._num_shuffle_parts)
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._start_idx * self._chunk_size)
if self._start_idx >= self._db_size: self.next_chunk()
self._end_idx = self._start_idx + self._chunk_size
self._end_idx = min(self._db_size, self._end_idx)
else:
self._start_idx = 0
self._end_idx = self._db_size
self.redirect(self._start_idx)
def next_record(self):
"""Step the cursor of records.
Returns
-------
None
"""
self._cur_idx += 1
self._db.next()
def next_chunk(self):
"""Step the cursor of shuffling chunks.
Returns
-------
None
"""
self._cur_chunk_idx += 1
if self._cur_chunk_idx >= self._num_shuffle_parts: self.reset()
else:
self._start_idx = self._part_idx * self._num_shuffle_parts + self._perm[self._cur_chunk_idx]
self._start_idx = self._start_idx * self._chunk_size
if self._start_idx >= self._db_size: self.next_chunk()
else:
self._end_idx = self._start_idx + self._chunk_size
self._end_idx = min(self._db_size, self._end_idx)
self.redirect(self._start_idx)
def run(self):
"""Start the process.
Returns
-------
None
"""
# fix seed
npr.seed(self._random_seed)
# init db
self._db = LMDB()
self._db.open(self._source)
self._db_size = int(self._db.get('size'))
self._db_zfill = int(self._db.get('zfill'))
self._epoch_size = int(self._db_size / self._num_parts + 1)
if self._use_instance_chunk:
self._chunk_size = 1
self._num_shuffle_parts = int(self._db_size / self._chunk_size / self._num_parts) + 1
else:
# search a optimal chunk size by chunks
if self._chunk_size == -1:
max_chunk_size = self._db._total_size / ((self._num_chunks * (1 << 20)))
min_chunk_size = 1
while min_chunk_size * 2 < max_chunk_size: min_chunk_size *= 2
self._chunk_size = min_chunk_size
self._num_shuffle_parts = int(math.ceil(self._db._total_size * 1.1 /
(self._num_parts * self._chunk_size << 20)))
self._chunk_size = int(self._db_size / self._num_shuffle_parts / self._num_parts + 1)
self._perm = np.arange(self._num_shuffle_parts)
# init env
self.reset()
# run
while True:
self.Q_out.put(self.element())
self.next_record()
if self._cur_idx >= self._end_idx:
if self._multiple_nodes or \
self._use_shuffle: self.next_chunk()
else: self.reset()
\ No newline at end of file
Dragon/python/dragon/vm/torch/utils/data/io/data_transformer.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>
#
# ------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
import numpy.random as npr
from multiprocessing import Process
import dragon.config as config
import dragon.vm.caffe.proto.caffe_pb2 as pb
try:
import cv2
except ImportError as e:
print('Failed to import cv2. Error: {0}'.format(str(e)))
class DataTransformer(Process):
"""DataTransformer is deployed to queue transformed images from `DataReader`_.
Nearly all common image augmentation methods are supported.
"""
def __init__(self, transform=None, color_space='RGB', pack=False, **kwargs):
"""Construct a ``DataTransformer``.
Parameters
----------
transform : lambda
The transforms.
color_space : str
The color space.
pack : boolean
Pack the images automatically.
"""
super(DataTransformer, self).__init__()
self.transform = transform
self.color_space = color_space
self.pack = pack
self._random_seed = config.GetRandomSeed()
self.Q_in = self.Q_out = None
self.daemon = True
def get(self, serialized):
"""Return image and labels from a serialized str.
Parameters
----------
serialized : str
The protobuf serialized str.
Returns
-------
tuple
The tuple image and labels.
"""
# decode
datum = pb.Datum()
datum.ParseFromString(serialized)
im = np.fromstring(datum.data, np.uint8)
if datum.encoded is True:
im = cv2.imdecode(im, -1)
else:
im = im.reshape((datum.height, datum.width, datum.channels))
if datum.channels == 3 and \
self.color_space == 'RGB': im = im[:, :, ::-1]
# labels
labels = []
if len(datum.labels) > 0: labels.extend(datum.labels)
else: labels.append(datum.label)
return self.transform(im), labels
def run(self):
"""Start the process.
Returns
-------
None
"""
npr.seed(self._random_seed)
while True:
serialized = self.Q_in.get()
im, label = self.get(serialized)
if len(im.shape) == 4 and not self.pack:
for ix in range(im.shape[0]):
self.Q_out.put((im[ix], label))
else:
if len(im.shape) == 3 and self.pack:
im = np.expand_dims(im, axis=0)
self.Q_out.put((im, label))
\ No newline at end of file
Dragon/python/dragon/vm/torch/utils/model_zoo.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>
#
# Codes are based on:
#
# <https://github.com/pytorch/pytorch/blob/master/torch/utils/model_zoo.py>
#
# ------------------------------------------------------------
import torch
import hashlib
import os
import re
import shutil
import sys
import tempfile
try:
from requests.utils import urlparse
import requests.get as urlopen
requests_available = True
except ImportError:
requests_available = False
if sys.version_info[0] == 2:
from urlparse import urlparse # noqa f811
from urllib2 import urlopen # noqa f811
else:
from urllib.request import urlopen
from urllib.parse import urlparse
try:
from tqdm import tqdm
except ImportError:
tqdm = None # defined below
# matches bfd8deac from resnet18-bfd8deac.pth
HASH_REGEX = re.compile(r'-([a-f0-9]*)\.')
def load_url(url, model_dir=None, map_location=None, progress=True):
r"""Loads the Torch serialized object at the given URL.
If the object is already present in `model_dir`, it's deserialized and
returned. The filename part of the URL should follow the naming convention
``filename-<sha256>.ext`` where ``<sha256>`` is the first eight or more
digits of the SHA256 hash of the contents of the file. The hash is used to
ensure unique names and to verify the contents of the file.
The default value of `model_dir` is ``$TORCH_HOME/models`` where
``$TORCH_HOME`` defaults to ``~/.torch``. The default directory can be
overriden with the ``$TORCH_MODEL_ZOO`` environment variable.
Args:
url (string): URL of the object to download
model_dir (string, optional): directory in which to save the object
map_location (optional): a function or a dict specifying how to remap storage locations (see torch.load)
progress (bool, optional): whether or not to display a progress bar to stderr
Example:
>>> state_dict = torch.utils.model_zoo.load_url('https://s3.amazonaws.com/pytorch/models/resnet18-5c106cde.pth')
"""
if model_dir is None:
torch_home = os.path.expanduser(os.getenv('TORCH_HOME', '~/.torch'))
model_dir = os.getenv('TORCH_MODEL_ZOO', os.path.join(torch_home, 'models'))
if not os.path.exists(model_dir):
os.makedirs(model_dir)
parts = urlparse(url)
filename = os.path.basename(parts.path)
cached_file = os.path.join(model_dir, filename)
if not os.path.exists(cached_file):
sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file))
hash_prefix = HASH_REGEX.search(filename).group(1)
_download_url_to_file(url, cached_file, hash_prefix, progress=progress)
return torch.load(cached_file, map_location=map_location)
def _download_url_to_file(url, dst, hash_prefix, progress):
u = urlopen(url)
if requests_available:
file_size = int(u.headers["Content-Length"])
u = u.raw
else:
meta = u.info()
if hasattr(meta, 'getheaders'):
file_size = int(meta.getheaders("Content-Length")[0])
else:
file_size = int(meta.get_all("Content-Length")[0])
f = tempfile.NamedTemporaryFile(delete=False)
try:
sha256 = hashlib.sha256()
with tqdm(total=file_size, disable=not progress) as pbar:
while True:
buffer = u.read(8192)
if len(buffer) == 0:
break
f.write(buffer)
sha256.update(buffer)
pbar.update(len(buffer))
f.close()
digest = sha256.hexdigest()
if digest[:len(hash_prefix)] != hash_prefix:
raise RuntimeError('invalid hash value (expected "{}", got "{}")'
.format(hash_prefix, digest))
shutil.move(f.name, dst)
finally:
f.close()
if os.path.exists(f.name):
os.remove(f.name)
if tqdm is None:
# fake tqdm if it's not installed
class tqdm(object):
def __init__(self, total, disable=False):
self.total = total
self.disable = disable
self.n = 0
def update(self, n):
if self.disable:
return
self.n += n
sys.stderr.write("\r{0:.1f}%".format(100 * self.n / float(self.total)))
sys.stderr.flush()
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if self.disable:
return
sys.stderr.write('\n')
\ No newline at end of file
Dragon/python/dragon/vm/torch/vision/__init__.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>
#
# ------------------------------------------------------------
\ No newline at end of file
Dragon/python/dragon/vm/torch/vision/models/__init__.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>
#
# ------------------------------------------------------------
from .alexnet import *
from .resnet import *
from .vgg import *
from .squeezenet import *
from .inception import *
\ No newline at end of file
Dragon/python/dragon/vm/torch/vision/models/alexnet.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>
#
# Codes are based on:
#
# <https://github.com/pytorch/vision/blob/master/torchvision/models/alexnet.py>
#
# ------------------------------------------------------------
"""Recommend hyper-parameters:
Nesterov-SGD, batch_size: 256, base_lr: 0.01, weight_decay: 0.0005
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import dragon.vm.torch.nn as nn
import dragon.vm.torch.utils.model_zoo as model_zoo
__all__ = ['AlexNet', 'alexnet']
model_urls = {
# Top1 = 0.5580, Top5 = 0.7858
'alexnet': 'http://dragon.seetatech.com/download/models/pytorch/imagenet/alexnet-owt-4df8aa71.pkl',
}
class AlexNet(nn.Module):
def __init__(self, num_classes=1000):
super(AlexNet, self).__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
nn.Conv2d(64, 192, kernel_size=5, padding=2),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
nn.Conv2d(192, 384, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(384, 256, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(256, 256, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
)
self.classifier = nn.Sequential(
nn.Dropout(),
nn.Linear(256 * 6 * 6, 4096),
nn.ReLU(inplace=True),
nn.Dropout(),
nn.Linear(4096, 4096),
nn.ReLU(inplace=True),
nn.Linear(4096, num_classes),
)
def forward(self, x):
x = self.features(x)
x = x.view(x.size(0), 256 * 6 * 6)
x = self.classifier(x)
return x
def alexnet(pretrained=False, **kwargs):
model = AlexNet(**kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['alexnet']))
return model
\ No newline at end of file
Dragon/python/dragon/vm/torch/vision/models/resnet.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>
#
# Codes are based on:
#
# <https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py>
#
# ------------------------------------------------------------
"""We add the zero-init-bn, comparing to the one in original model zoo.
For more about zero-init, See,
"Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour"
<http://arxiv.org/abs/1706.02677>
Recommend hyper-parameters:
Nesterov-SGD, batch_size: 256, base_lr: 0.1, weight_decay: 0.0001
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import dragon.vm.torch.nn as nn
import dragon.vm.torch.utils.model_zoo as model_zoo
__all__ = [
'ResNet', 'resnet18', 'resnet34',
'resnet50', 'resnet101', 'resnet152'
]
model_urls = {
# Top1 = 0.6866, Top5 = 0.8855
'resnet18': 'http://dragon.seetatech.com/download/models/pytorch/imagenet/resnet18-5c106cde.pkl',
# Top1 = 0.7268, Top5 = 0.9102
'resnet34': 'http://dragon.seetatech.com/download/models/pytorch/imagenet/resnet34-333f7ec4.pkl',
# Top1 = 0.7506, Top5 = 0.9238
'resnet50': 'http://dragon.seetatech.com/download/models/pytorch/imagenet/resnet50-19c8e357.pkl',
# Top1 = 0.7666, Top5 = 0.9327
'resnet101': 'http://dragon.seetatech.com/download/models/pytorch/imagenet/resnet101-5d3b4d8f.pkl',
# Top1 = 0.7760, Top5 = 0.9374
'resnet152': 'http://dragon.seetatech.com/download/models/pytorch/imagenet/resnet152-b121ed2d.pkl',
}
def conv3x3(in_planes, out_planes, stride=1):
"""3x3 convolution with padding"""
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
padding=1, bias=False)
class BasicBlock(nn.Module):
expansion = 1
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(BasicBlock, self).__init__()
self.conv1 = conv3x3(inplanes, planes, stride)
self.bn1 = nn.BatchNorm2d(planes)
self.relu = nn.ReLU(inplace=True)
self.conv2 = conv3x3(planes, planes)
self.bn2 = nn.BatchNorm2d(planes)
self.bn2.zero_init = True
self.downsample = downsample
self.stride = stride
def forward(self, x):
residual = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
if self.downsample is not None:
residual = self.downsample(x)
out += residual
out = self.relu(out)
return out
class Bottleneck(nn.Module):
expansion = 4
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(Bottleneck, self).__init__()
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
self.bn1 = nn.BatchNorm2d(planes)
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(planes)
self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)
self.bn3 = nn.BatchNorm2d(planes * self.expansion)
self.bn3.zero_init = True
self.relu = nn.ReLU(inplace=True)
self.downsample = downsample
self.stride = stride
def forward(self, x):
residual = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
out = self.relu(out)
out = self.conv3(out)
out = self.bn3(out)
if self.downsample is not None:
residual = self.downsample(x)
out += residual
out = self.relu(out)
return out
class ResNet(nn.Module):
def __init__(self, block, layers, num_classes=1000):
self.inplanes = 64
super(ResNet, self).__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(block, 64, layers[0])
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
self.avgpool = nn.AvgPool2d(7, stride=1)
self.fc = nn.Linear(512 * block.expansion, num_classes)
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
elif isinstance(m, nn.BatchNorm2d):
if hasattr(m, 'zero_init'):
nn.init.constant_(m.weight, 0)
else:
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
def _make_layer(self, block, planes, blocks, stride=1):
downsample = None
if stride != 1 or self.inplanes != planes * block.expansion:
downsample = nn.Sequential(
nn.Conv2d(self.inplanes, planes * block.expansion,
kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(planes * block.expansion),
)
layers = []
layers.append(block(self.inplanes, planes, stride, downsample))
self.inplanes = planes * block.expansion
for i in range(1, blocks):
layers.append(block(self.inplanes, planes))
return nn.Sequential(*layers)
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.avgpool(x)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
def resnet18(pretrained=False, **kwargs):
model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['resnet18']))
return model
def resnet34(pretrained=False, **kwargs):
model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['resnet34']))
return model
def resnet50(pretrained=False, **kwargs):
model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
return model
def resnet101(pretrained=False, **kwargs):
model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))
return model
def resnet152(pretrained=False, **kwargs):
model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))
return model
\ No newline at end of file
Dragon/python/dragon/vm/torch/vision/models/squeezenet.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>
#
# Codes are based on:
#
# <https://github.com/pytorch/vision/blob/master/torchvision/models/squeezenet.py>
#
# ------------------------------------------------------------
"""Recommend hyper-parameters:
Nesterov-SGD, batch_size: 512, base_lr: 0.04, weight_decay: 0.0002
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import dragon.vm.torch as torch
import dragon.vm.torch.nn as nn
import dragon.vm.torch.nn.init as init
import dragon.vm.torch.utils.model_zoo as model_zoo
__all__ = ['SqueezeNet', 'squeezenet1_0', 'squeezenet1_1']
model_urls = {
# Top1 = 0.5674, Top5 = 0.7971
'squeezenet1_0': 'http://dragon.seetatech.com/download/models/pytorch/imagenet/squeezenet1_0-a815701f.pkl',
# Top1 = 0.5671, Top5 = 0.7955
'squeezenet1_1': 'http://dragon.seetatech.com/download/models/pytorch/imagenet/squeezenet1_1-f364aa15.pkl',
}
class Fire(nn.Module):
def __init__(self, inplanes, squeeze_planes,
expand1x1_planes, expand3x3_planes):
super(Fire, self).__init__()
self.inplanes = inplanes
self.squeeze = nn.Conv2d(inplanes, squeeze_planes, kernel_size=1)
self.squeeze_activation = nn.ReLU(inplace=True)
self.expand1x1 = nn.Conv2d(squeeze_planes, expand1x1_planes,
kernel_size=1)
self.expand1x1_activation = nn.ReLU(inplace=True)
self.expand3x3 = nn.Conv2d(squeeze_planes, expand3x3_planes,
kernel_size=3, padding=1)
self.expand3x3_activation = nn.ReLU(inplace=True)
def forward(self, x):
x = self.squeeze_activation(self.squeeze(x))
return torch.cat([
self.expand1x1_activation(self.expand1x1(x)),
self.expand3x3_activation(self.expand3x3(x))
], 1)
class SqueezeNet(nn.Module):
def __init__(self, version=1.0, num_classes=1000):
super(SqueezeNet, self).__init__()
if version not in [1.0, 1.1]:
raise ValueError("Unsupported SqueezeNet version {version}:"
"1.0 or 1.1 expected".format(version=version))
self.num_classes = num_classes
if version == 1.0:
self.features = nn.Sequential(
nn.Conv2d(3, 96, kernel_size=7, stride=2),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
Fire(96, 16, 64, 64),
Fire(128, 16, 64, 64),
Fire(128, 32, 128, 128),
nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
Fire(256, 32, 128, 128),
Fire(256, 48, 192, 192),
Fire(384, 48, 192, 192),
Fire(384, 64, 256, 256),
nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
Fire(512, 64, 256, 256),
)
else:
self.features = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=3, stride=2),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
Fire(64, 16, 64, 64),
Fire(128, 16, 64, 64),
nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
Fire(128, 32, 128, 128),
Fire(256, 32, 128, 128),
nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
Fire(256, 48, 192, 192),
Fire(384, 48, 192, 192),
Fire(384, 64, 256, 256),
Fire(512, 64, 256, 256),
)
# Final convolution is initialized differently form the rest
final_conv = nn.Conv2d(512, self.num_classes, kernel_size=1)
self.classifier = nn.Sequential(
nn.Dropout(p=0.5),
final_conv,
nn.ReLU(inplace=True),
nn.AvgPool2d(13, stride=1)
)
for m in self.modules():
if isinstance(m, nn.Conv2d):
if m is final_conv:
init.normal_(m.weight, mean=0.0, std=0.01)
else:
init.kaiming_uniform_(m.weight)
if m.bias is not None:
init.constant_(m.bias, 0)
def forward(self, x):
x = self.features(x)
x = self.classifier(x)
return x.view(x.size(0), self.num_classes)
def squeezenet1_0(pretrained=False, **kwargs):
model = SqueezeNet(version=1.0, **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['squeezenet1_0']))
return model
def squeezenet1_1(pretrained=False, **kwargs):
model = SqueezeNet(version=1.1, **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['squeezenet1_1']))
return model
\ No newline at end of file
Dragon/python/dragon/vm/torch/vision/models/vgg.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>
#
# Codes are based on:
#
# <https://github.com/pytorch/vision/blob/master/torchvision/models/vgg.py>
#
# ------------------------------------------------------------
"""Recommend hyper-parameters:
Nesterov-SGD, batch_size: 256, base_lr: 0.01, weight_decay: 0.0005
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import dragon.vm.torch.nn as nn
import dragon.vm.torch.utils.model_zoo as model_zoo
__all__ = [
'VGG', 'vgg11', 'vgg11_bn', 'vgg13', 'vgg13_bn', 'vgg16', 'vgg16_bn',
'vgg19_bn', 'vgg19',
]
model_urls = {
# Top1 = 0.6776, Top5 = 0.8799
'vgg11': 'http://dragon.seetatech.com/download/models/pytorch/imagenet/vgg11-bbd30ac9.pkl',
# Top1 = 0.6867, Top5 = 0.8855
'vgg13': 'http://dragon.seetatech.com/download/models/pytorch/imagenet/vgg13-c768596a.pkl',
# Top1 = 0.7057, Top5 = 0.8974
'vgg16': 'http://dragon.seetatech.com/download/models/pytorch/imagenet/vgg16-397923af.pkl',
# Top1 = 0.7147, Top5 = 0.9017
'vgg19': 'http://dragon.seetatech.com/download/models/pytorch/imagenet/vgg19-dcbb9e9d.pkl',
# Top1 = 0.6939, Top5 = 0.8918
'vgg11_bn': 'http://dragon.seetatech.com/download/models/pytorch/imagenet/vgg11_bn-6002323d.pkl',
# Top1 = 0.7020, Top5 = 0.8957
'vgg13_bn': 'http://dragon.seetatech.com/download/models/pytorch/imagenet/vgg13_bn-abd245e5.pkl',
# Top1 = 0.7225, Top5 = 0.9086
'vgg16_bn': 'http://dragon.seetatech.com/download/models/pytorch/imagenet/vgg16_bn-6c64b313.pkl',
# Top1 = 0.7319, Top5 = 0.9133
'vgg19_bn': 'http://dragon.seetatech.com/download/models/pytorch/imagenet/vgg19_bn-c79401a0.pkl',
}
class VGG(nn.Module):
def __init__(self, features, num_classes=1000, init_weights=True):
super(VGG, self).__init__()
self.features = features
self.classifier = nn.Sequential(
nn.Linear(512 * 7 * 7, 4096),
nn.ReLU(True),
nn.Dropout(),
nn.Linear(4096, 4096),
nn.ReLU(True),
nn.Dropout(),
nn.Linear(4096, num_classes),
)
if init_weights:
self._initialize_weights()
def forward(self, x):
x = self.features(x)
x = x.view(x.size(0), -1)
x = self.classifier(x)
return x
def _initialize_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
if m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.BatchNorm2d):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.Linear):
nn.init.normal_(m.weight, 0, 0.01)
nn.init.constant_(m.bias, 0)
def make_layers(cfg, batch_norm=False):
layers = []
in_channels = 3
for v in cfg:
if v == 'M':
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
else:
conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
if batch_norm:
layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
else:
layers += [conv2d, nn.ReLU(inplace=True)]
in_channels = v
return nn.Sequential(*layers)
cfg = {
'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
}
def vgg11(pretrained=False, **kwargs):
if pretrained:
kwargs['init_weights'] = False
model = VGG(make_layers(cfg['A']), **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['vgg11']))
return model
def vgg11_bn(pretrained=False, **kwargs):
if pretrained:
kwargs['init_weights'] = False
model = VGG(make_layers(cfg['A'], batch_norm=True), **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['vgg11_bn']))
return model
def vgg13(pretrained=False, **kwargs):
if pretrained:
kwargs['init_weights'] = False
model = VGG(make_layers(cfg['B']), **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['vgg13']))
return model
def vgg13_bn(pretrained=False, **kwargs):
if pretrained:
kwargs['init_weights'] = False
model = VGG(make_layers(cfg['B'], batch_norm=True), **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['vgg13_bn']))
return model
def vgg16(pretrained=False, **kwargs):
if pretrained:
kwargs['init_weights'] = False
model = VGG(make_layers(cfg['D']), **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['vgg16']))
return model
def vgg16_bn(pretrained=False, **kwargs):
if pretrained:
kwargs['init_weights'] = False
model = VGG(make_layers(cfg['D'], batch_norm=True), **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['vgg16_bn']))
return model
def vgg19(pretrained=False, **kwargs):
if pretrained:
kwargs['init_weights'] = False
model = VGG(make_layers(cfg['E']), **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['vgg19']))
return model
def vgg19_bn(pretrained=False, **kwargs):
if pretrained:
kwargs['init_weights'] = False
model = VGG(make_layers(cfg['E'], batch_norm=True), **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['vgg19_bn']))
return model
\ No newline at end of file
Dragon/python/dragon/vm/torch/vision/transforms/__init__.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>
#
# ------------------------------------------------------------
\ No newline at end of file
Dragon/python/dragon/vm/torch/vision/transforms/functional.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>
#
#
# Codes are based on:
#
# <https://github.com/pytorch/vision/blob/master/torchvision/transforms/functional.py>
#
# ------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import sys
import cv2
from PIL import Image, ImageEnhance
try:
import accimage
except ImportError:
accimage = None
import numpy as np
import numbers
import collections
import warnings
def _is_pil_image(img):
if accimage is not None:
return isinstance(img, (Image.Image, accimage.Image))
else:
return isinstance(img, Image.Image)
def _is_numpy_image(img):
return isinstance(img, np.ndarray) and (img.ndim in {2, 3, 4})
def to_array(pic, normalize=True):
"""Convert a ``PIL Image`` to ``numpy.ndarray``
"""
if not _is_numpy_image(pic):
raise TypeError('pic should be numpy Image. Got {}'.format(type(pic)))
if normalize: return pic.astype(np.float32) / 255
else: return pic.astype(np.float32)
def normalize(array, mean, std):
if not _is_numpy_image(array):
raise TypeError('array is not a numpy image.')
array = (array - mean) / std
array = array.astype(np.float32)
if len(array.shape) == 3: return array.transpose(2, 0, 1)
elif len(array.shape) == 4: return array.transpose(0, 3, 1, 2)
else: raise ValueError('excepted a 3d or 4d array.')
_pil_interpolation_to_ocv = {
Image.NEAREST: cv2.INTER_NEAREST,
Image.BILINEAR: cv2.INTER_LINEAR,
Image.BICUBIC: cv2.INTER_CUBIC,
Image.LANCZOS: cv2.INTER_LANCZOS4
}
def resize(img, size, interpolation=Image.BILINEAR):
if not _is_numpy_image(img):
raise TypeError('img should be numpy Image. Got {}'.format(type(img)))
if not (isinstance(size, int) or (isinstance(size, collections.Iterable) and len(size) == 2)):
raise TypeError('Got inappropriate size arg: {}'.format(size))
if isinstance(size, int):
h, w = img.shape[0:2]
if (w <= h and w == size) or (h <= w and h == size):
return img
if w < h:
ow = size
oh = int(size * h / w)
else:
oh = size
ow = int(size * w / h)
else: oh, ow = size
if sys.version_info >= (3, 0):
return cv2.resize(img, dsize=(ow, oh),
interpolation=_pil_interpolation_to_ocv[interpolation])
else:
# Fuck Fuck Fuck opencv-python2, it always has a BUG
# that leads to duplicate cuDA handles created at gpu:0
img = Image.fromarray(img)
img.resize((ow, oh), interpolation)
return np.array(img)
def scale(*args, **kwargs):
warnings.warn("The use of the transforms.Scale transform is deprecated, " +
"please use transforms.Resize instead.")
return resize(*args, **kwargs)
def pad(img, padding, fill=0, padding_mode='constant'):
if not _is_numpy_image(img):
raise TypeError('img should be numpy Image. Got {}'.format(type(img)))
if not isinstance(padding, (numbers.Number, tuple)):
raise TypeError('Got inappropriate padding arg')
if not isinstance(fill, (numbers.Number, str, tuple)):
raise TypeError('Got inappropriate fill arg')
if not isinstance(padding_mode, str):
raise TypeError('Got inappropriate padding_mode arg')
if isinstance(padding, collections.Sequence) and len(padding) not in [2, 4]:
raise ValueError("Padding must be an int or a 2, or 4 element tuple, not a " +
"{} element tuple".format(len(padding)))
assert padding_mode in ['constant', 'edge', 'reflect', 'symmetric'], \
'Padding mode should be either constant, edge, reflect or symmetric'
if isinstance(padding, int):
pad_left = pad_right = pad_top = pad_bottom = padding
if isinstance(padding, collections.Sequence) and len(padding) == 2:
pad_left = pad_right = padding[0]
pad_top = pad_bottom = padding[1]
if isinstance(padding, collections.Sequence) and len(padding) == 4:
pad_left = padding[0]
pad_top = padding[1]
pad_right = padding[2]
pad_bottom = padding[3]
# RGB image
if len(img.shape) == 3:
img = np.pad(img, ((pad_top, pad_bottom), (pad_left, pad_right), (0, 0)), padding_mode)
# Grayscale image
if len(img.shape) == 2:
img = np.pad(img, ((pad_top, pad_bottom), (pad_left, pad_right)), padding_mode)
return img
def crop(img, i, j, h, w, copy=False):
if not _is_numpy_image(img):
raise TypeError('img should be numpy Image. Got {}'.format(type(img)))
if copy: return img[i : i + h, j : j + w].copy()
return img[i : i + h, j : j + w]
def center_crop(img, output_size):
if isinstance(output_size, numbers.Number):
output_size = (int(output_size), int(output_size))
h, w = img.shape[0:2]
th, tw = output_size
i = int(round((h - th) / 2.))
j = int(round((w - tw) / 2.))
return crop(img, i, j, th, tw)
def resized_crop(img, i, j, h, w, size, interpolation=Image.BILINEAR):
assert _is_numpy_image(img), 'img should be numpy Image'
img = crop(img, i, j, h, w)
img = resize(img, size, interpolation)
return img
def hflip(img):
if not _is_numpy_image(img):
raise TypeError('img should be numpy Image. Got {}'.format(type(img)))
return img[:, ::-1]
def vflip(img):
if not _is_numpy_image(img):
raise TypeError('img should be numpy Image. Got {}'.format(type(img)))
return img[::-1, :]
def five_crop(img, size):
if isinstance(size, numbers.Number):
size = (int(size), int(size))
else:
assert len(size) == 2, "Please provide only two dimensions (h, w) for size."
h, w = img.shape[0:2]
crop_h, crop_w = size
if crop_w > w or crop_h > h:
raise ValueError("Requested crop size {} is bigger than input size {}".format(size, (h, w)))
tl = crop(img, 0, 0, crop_h, crop_w, True)
tr = crop(img, 0, w - crop_w, crop_h, crop_w, True)
bl = crop(img, h - crop_h, 0, crop_h, crop_w, True)
br = crop(img, h - crop_h, w - crop_w, crop_h, crop_w, True)
center = center_crop(img, (crop_h, crop_w)).copy()
return np.stack((tl, tr, bl, br, center), axis=0)
def ten_crop(img, size, vertical_flip=False):
if isinstance(size, numbers.Number):
size = (int(size), int(size))
else:
assert len(size) == 2, "Please provide only two dimensions (h, w) for size."
first_five = five_crop(img, size)
if vertical_flip:
img = vflip(img)
else:
img = hflip(img)
second_five = five_crop(img, size)
return np.concatenate([first_five, second_five], axis=0)
def adjust_brightness(img, brightness_factor):
if not _is_pil_image(img):
raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
enhancer = ImageEnhance.Brightness(img)
img = enhancer.enhance(brightness_factor)
return img
def adjust_contrast(img, contrast_factor):
if not _is_pil_image(img):
raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
enhancer = ImageEnhance.Contrast(img)
img = enhancer.enhance(contrast_factor)
return img
def adjust_saturation(img, saturation_factor):
if not _is_pil_image(img):
raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
enhancer = ImageEnhance.Color(img)
img = enhancer.enhance(saturation_factor)
return img
def adjust_hue(img, hue_factor):
if not(-0.5 <= hue_factor <= 0.5):
raise ValueError('hue_factor is not in [-0.5, 0.5].'.format(hue_factor))
if not _is_pil_image(img):
raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
input_mode = img.mode
if input_mode in {'L', '1', 'I', 'F'}:
return img
h, s, v = img.convert('HSV').split()
np_h = np.array(h, dtype=np.uint8)
# uint8 addition take cares of rotation across boundaries
with np.errstate(over='ignore'):
np_h += np.uint8(hue_factor * 255)
h = Image.fromarray(np_h, 'L')
img = Image.merge('HSV', (h, s, v)).convert(input_mode)
return img
\ 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!