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 3356 additions and 343 deletions
Docker/ubuntu-16.04-cpu-openblas/Dockerfile
...@@ -8,6 +8,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ ...@@ -8,6 +8,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
unzip \ unzip \
ssh \ ssh \
vim \ vim \
libnuma-dev \
libprotobuf-dev \ libprotobuf-dev \
protobuf-compiler \ protobuf-compiler \
libopenblas-dev \ libopenblas-dev \
...@@ -17,7 +18,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ ...@@ -17,7 +18,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
python3-tk \ python3-tk \
&& rm -rf /var/lib/apt/lists/* && 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 \ pip3 install --no-cache-dir -i https://pypi.tuna.tsinghua.edu.cn/simple \
numpy \ numpy \
protobuf \ protobuf \
......
Docker/ubuntu-16.04-cuda9.0-cudnn7/Dockerfile
...@@ -8,6 +8,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ ...@@ -8,6 +8,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
unzip \ unzip \
ssh \ ssh \
vim \ vim \
libnuma-dev \
libprotobuf-dev \ libprotobuf-dev \
protobuf-compiler \ protobuf-compiler \
libopenblas-dev \ libopenblas-dev \
...@@ -19,7 +20,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ ...@@ -19,7 +20,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
python3-tk \ python3-tk \
&& rm -rf /var/lib/apt/lists/* && 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 \ pip3 install --no-cache-dir -i https://pypi.tuna.tsinghua.edu.cn/simple \
numpy \ numpy \
protobuf \ protobuf \
......
Dragon/include/core/operator.h
...@@ -68,6 +68,7 @@ class OperatorBase { ...@@ -68,6 +68,7 @@ class OperatorBase {
inline const OperatorDef& op_def() const { return op_def_; } inline const OperatorDef& op_def() const { return op_def_; }
inline string DebugString() const { return op_def_.DebugString(); } inline string DebugString() const { return op_def_.DebugString(); }
string DTypeHelper(const Tensor& tensor, const Set<string>& dtypes) const; string DTypeHelper(const Tensor& tensor, const Set<string>& dtypes) const;
string DTypeHelper(const string& dtype, const Set<string>& dtypes) const;
protected: protected:
string phase_, anchor_; string phase_, anchor_;
......
Dragon/include/core/tensor.h
...@@ -206,9 +206,8 @@ class Tensor { ...@@ -206,9 +206,8 @@ class Tensor {
template <class DstCTX, class SrcCTX> template <class DstCTX, class SrcCTX>
inline void Copy(const Tensor& other) { inline void Copy(const Tensor& other) {
CHECK_EQ(size_, other.size_); CHECK_EQ(size_, other.size_);
meta_ = other.meta_;
auto* src = other.template raw_data<SrcCTX>(); 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 (dst == src) return;
if (TypeMeta::Id<DstCTX>() == TypeMeta::Id<CPUContext>()) { if (TypeMeta::Id<DstCTX>() == TypeMeta::Id<CPUContext>()) {
CPUContext::Memcpy<DstCTX, SrcCTX>(nbytes(), dst, src); CPUContext::Memcpy<DstCTX, SrcCTX>(nbytes(), dst, src);
......
Dragon/include/operators/activation/dropout_op.h
...@@ -33,7 +33,6 @@ class DropoutOp final : public Operator<Context> { ...@@ -33,7 +33,6 @@ class DropoutOp final : public Operator<Context> {
protected: protected:
DECLARE_ARGUMENT_WITH_DESC(float, prob); DECLARE_ARGUMENT_WITH_DESC(float, prob);
bool use_scale; bool use_scale;
Tensor* mask;
}; };
template <class Context> template <class Context>
......
Dragon/include/operators/cast/float2half_op.h → Dragon/include/operators/misc/astype_op.h
...@@ -9,21 +9,29 @@ ...@@ -9,21 +9,29 @@
// //
// ------------------------------------------------------------ // ------------------------------------------------------------
#ifndef DRAGON_OPERATORS_CAST_FLOAT2HALF_OP_H_ #ifndef DRAGON_OPERATORS_MISC_ASTYPE_OP_H_
#define DRAGON_OPERATORS_CAST_FLOAT2HALF_OP_H_ #define DRAGON_OPERATORS_MISC_ASTYPE_OP_H_
#include "core/operator.h" #include "core/operator.h"
namespace dragon { namespace dragon {
template <class Context> template <class Context>
class FloatToHalfOp final : public Operator<Context> { class AsTypeOp final : public Operator<Context> {
public: 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); USE_OPERATOR_FUNCTIONS(Context);
void RunOnDevice() override; void RunOnDevice() override;
protected:
string dtype;
bool inplace;
}; };
} // namespace dragon } // namespace dragon
#endif // DRAGON_OPERATORS_CAST_FLOAT2HALF_OP_H_ #endif // DRAGON_OPERATORS_MISC_ASTYPE_OP_H_
\ No newline at end of file \ 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 @@ ...@@ -9,19 +9,20 @@
// //
// ------------------------------------------------------------- // -------------------------------------------------------------
#ifndef DRAGON_OPERATORS_NDARRAY_ARGMAX_OP_H_ #ifndef DRAGON_OPERATORS_NDARRAY_ARGREDUCE_OP_H_
#define DRAGON_OPERATORS_NDARRAY_ARGMAX_OP_H_ #define DRAGON_OPERATORS_NDARRAY_ARGREDUCE_OP_H_
#include "core/operator.h" #include "core/operator.h"
namespace dragon { namespace dragon {
template <class Context> template <class Context>
class ArgmaxOp final : public Operator<Context> { class ArgReduceOp final : public Operator<Context> {
public: public:
ArgmaxOp(const OperatorDef& op_def, Workspace* ws) ArgReduceOp(const OperatorDef& op_def, Workspace* ws)
: Operator<Context>(op_def, ws), : Operator<Context>(op_def, ws),
axis(OperatorBase::GetSingleArg<int>("axis", -1)), axis(OperatorBase::GetSingleArg<int>("axis", -1)),
operation(OperatorBase::GetSingleArg<string>("operation", "NONE")),
keep_dims(OperatorBase::GetSingleArg<bool>("keep_dims", false)), keep_dims(OperatorBase::GetSingleArg<bool>("keep_dims", false)),
top_k(OperatorBase::GetSingleArg<int>("top_k", 1)) {} top_k(OperatorBase::GetSingleArg<int>("top_k", 1)) {}
USE_OPERATOR_FUNCTIONS(Context); USE_OPERATOR_FUNCTIONS(Context);
...@@ -30,10 +31,11 @@ class ArgmaxOp final : public Operator<Context> { ...@@ -30,10 +31,11 @@ class ArgmaxOp final : public Operator<Context> {
template <typename T> void RunWithType(); template <typename T> void RunWithType();
protected: protected:
TIndex axis, axis_dim, top_k, count, inner_dim;
bool keep_dims; bool keep_dims;
string operation;
TIndex axis, axis_dim, top_k, count, inner_dim;
}; };
} // namespace dragon } // namespace dragon
#endif // DRAGON_OPERATORS_NDARRAY_ARGMAX_OP_H_ #endif // DRAGON_OPERATORS_NDARRAY_ARGREDUCE_OP_H_
\ No newline at end of file \ No newline at end of file
Dragon/include/operators/ndarray/concat_op.h
...@@ -21,15 +21,14 @@ class ConcatOp : public Operator<Context> { ...@@ -21,15 +21,14 @@ class ConcatOp : public Operator<Context> {
public: public:
ConcatOp(const OperatorDef& op_def, Workspace* ws) ConcatOp(const OperatorDef& op_def, Workspace* ws)
: Operator<Context>(op_def, ws), : Operator<Context>(op_def, ws),
axis(OperatorBase::GetSingleArg<int>("axis", 1)), axis(OperatorBase::GetSingleArg<int>("axis", 1)) {}
nin(OperatorBase::GetSingleArg<int>("num_input", 1)) {}
USE_OPERATOR_FUNCTIONS(Context); USE_OPERATOR_FUNCTIONS(Context);
void RunOnDevice() override; void RunOnDevice() override;
template <typename T> void RunWithType(); template <typename T> void RunWithType();
protected: 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; TIndex x_offset, y_offset, concat_offset;
vector<TIndex> concat_dims; vector<TIndex> concat_dims;
}; };
...@@ -39,15 +38,14 @@ class ConcatGradientOp : public Operator<Context> { ...@@ -39,15 +38,14 @@ class ConcatGradientOp : public Operator<Context> {
public: public:
ConcatGradientOp(const OperatorDef& op_def, Workspace* ws) ConcatGradientOp(const OperatorDef& op_def, Workspace* ws)
: Operator<Context>(op_def, ws), : Operator<Context>(op_def, ws),
axis(OperatorBase::GetSingleArg<int>("axis", 1)), axis(OperatorBase::GetSingleArg<int>("axis", 1)) {}
nin(OperatorBase::GetSingleArg<int>("num_input", 1)) {}
USE_OPERATOR_FUNCTIONS(Context); USE_OPERATOR_FUNCTIONS(Context);
void RunOnDevice() override; void RunOnDevice() override;
template <typename T> void RunWithType(); template <typename T> void RunWithType();
protected: 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; TIndex x_offset, y_offset, concat_offset;
vector<TIndex> concat_dims; vector<TIndex> concat_dims;
}; };
......
Dragon/include/operators/norm/batch_norm_op.h
...@@ -12,6 +12,8 @@ ...@@ -12,6 +12,8 @@
#ifndef DRAGON_OPERATORS_NORM_BATCH_NORM_OP_H_ #ifndef DRAGON_OPERATORS_NORM_BATCH_NORM_OP_H_
#define DRAGON_OPERATORS_NORM_BATCH_NORM_OP_H_ #define DRAGON_OPERATORS_NORM_BATCH_NORM_OP_H_
#include <cfloat>
#include "core/operator.h" #include "core/operator.h"
namespace dragon { namespace dragon {
...@@ -22,8 +24,8 @@ class BatchNormOp : public Operator<Context> { ...@@ -22,8 +24,8 @@ class BatchNormOp : public Operator<Context> {
BatchNormOp(const OperatorDef& op_def, Workspace* ws) BatchNormOp(const OperatorDef& op_def, Workspace* ws)
: Operator<Context>(op_def, ws), : Operator<Context>(op_def, ws),
axis(OperatorBase::GetSingleArg<int>("axis", -1)), axis(OperatorBase::GetSingleArg<int>("axis", -1)),
momentum(OperatorBase::GetSingleArg<float>("momentum", float(0.9))), momentum(OperatorBase::GetSingleArg<float>("momentum", 0.9f)),
eps(OperatorBase::GetSingleArg<float>("eps", float(1e-3))), eps(OperatorBase::GetSingleArg<float>("eps", 1e-3f)),
use_stats(OperatorBase::GetSingleArg<int>("use_stats", -1)), use_stats(OperatorBase::GetSingleArg<int>("use_stats", -1)),
mode(OperatorBase::GetSingleArg<string>("mode", "DEFAULT")) { mode(OperatorBase::GetSingleArg<string>("mode", "DEFAULT")) {
if (axis != -1) if (axis != -1)
...@@ -84,8 +86,8 @@ class FusedBatchNormOp : public Operator<Context> { ...@@ -84,8 +86,8 @@ class FusedBatchNormOp : public Operator<Context> {
FusedBatchNormOp(const OperatorDef& op_def, Workspace* ws) FusedBatchNormOp(const OperatorDef& op_def, Workspace* ws)
: Operator<Context>(op_def, ws), : Operator<Context>(op_def, ws),
axis(OperatorBase::GetSingleArg<int>("axis", -1)), axis(OperatorBase::GetSingleArg<int>("axis", -1)),
momentum(OperatorBase::GetSingleArg<float>("momentum", float(0.9))), momentum(OperatorBase::GetSingleArg<float>("momentum", 0.9f)),
eps(OperatorBase::GetSingleArg<float>("eps", float(1e-3))), eps(OperatorBase::GetSingleArg<float>("eps", 1e-3f)),
use_stats(OperatorBase::GetSingleArg<int>("use_stats", -1)) {} use_stats(OperatorBase::GetSingleArg<int>("use_stats", -1)) {}
USE_OPERATOR_FUNCTIONS(Context); USE_OPERATOR_FUNCTIONS(Context);
...@@ -112,7 +114,7 @@ class FusedBatchNormGradientOp : public Operator<Context> { ...@@ -112,7 +114,7 @@ class FusedBatchNormGradientOp : public Operator<Context> {
FusedBatchNormGradientOp(const OperatorDef& op_def, Workspace* ws) FusedBatchNormGradientOp(const OperatorDef& op_def, Workspace* ws)
: Operator<Context>(op_def, ws), : Operator<Context>(op_def, ws),
axis(OperatorBase::GetSingleArg<int>("axis", -1)), 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_stats(OperatorBase::GetSingleArg<int>("use_stats", -1)) {}
USE_OPERATOR_FUNCTIONS(Context); USE_OPERATOR_FUNCTIONS(Context);
...@@ -143,11 +145,16 @@ template <class Context> ...@@ -143,11 +145,16 @@ template <class Context>
class CuDNNBatchNormOp final : public FusedBatchNormOp<Context> { class CuDNNBatchNormOp final : public FusedBatchNormOp<Context> {
public: public:
CuDNNBatchNormOp(const OperatorDef& op_def, Workspace* ws) 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(&input_desc));
CUDNN_CHECK(cudnnCreateTensorDescriptor(&output_desc)); CUDNN_CHECK(cudnnCreateTensorDescriptor(&output_desc));
CUDNN_CHECK(cudnnCreateTensorDescriptor(&bn_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); USE_OPERATOR_FUNCTIONS(Context);
...@@ -163,6 +170,7 @@ class CuDNNBatchNormOp final : public FusedBatchNormOp<Context> { ...@@ -163,6 +170,7 @@ class CuDNNBatchNormOp final : public FusedBatchNormOp<Context> {
template <typename T> void RunWithType(); template <typename T> void RunWithType();
protected: protected:
double eps64;
cudnnTensorDescriptor_t input_desc, output_desc, bn_desc; cudnnTensorDescriptor_t input_desc, output_desc, bn_desc;
cudnnBatchNormMode_t bn_mode; cudnnBatchNormMode_t bn_mode;
TIndex N, C; TIndex N, C;
...@@ -174,11 +182,16 @@ template <class Context> ...@@ -174,11 +182,16 @@ template <class Context>
class CuDNNBatchNormGradientOp final : public FusedBatchNormGradientOp<Context> { class CuDNNBatchNormGradientOp final : public FusedBatchNormGradientOp<Context> {
public: public:
CuDNNBatchNormGradientOp(const OperatorDef& op_def, Workspace* ws) 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(&input_desc));
CUDNN_CHECK(cudnnCreateTensorDescriptor(&output_desc)); CUDNN_CHECK(cudnnCreateTensorDescriptor(&output_desc));
CUDNN_CHECK(cudnnCreateTensorDescriptor(&bn_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); USE_OPERATOR_FUNCTIONS(Context);
...@@ -195,6 +208,7 @@ class CuDNNBatchNormGradientOp final : public FusedBatchNormGradientOp<Context> ...@@ -195,6 +208,7 @@ class CuDNNBatchNormGradientOp final : public FusedBatchNormGradientOp<Context>
template <typename T> void InferenceRunWithType(); template <typename T> void InferenceRunWithType();
protected: protected:
double eps64;
cudnnTensorDescriptor_t input_desc, output_desc, bn_desc; cudnnTensorDescriptor_t input_desc, output_desc, bn_desc;
cudnnBatchNormMode_t bn_mode; cudnnBatchNormMode_t bn_mode;
TIndex N, C, S, NC, NS; TIndex N, C, S, NC, NS;
......
Dragon/include/operators/norm/batch_renorm_op.h
...@@ -22,13 +22,13 @@ class BatchRenormOp : public Operator<Context> { ...@@ -22,13 +22,13 @@ class BatchRenormOp : public Operator<Context> {
BatchRenormOp(const OperatorDef& op_def, Workspace* ws) BatchRenormOp(const OperatorDef& op_def, Workspace* ws)
: Operator<Context>(op_def, ws), : Operator<Context>(op_def, ws),
axis(OperatorBase::GetSingleArg<int>("axis", -1)), axis(OperatorBase::GetSingleArg<int>("axis", -1)),
momentum(OperatorBase::GetSingleArg<float>("momentum", float(0.9))), momentum(OperatorBase::GetSingleArg<float>("momentum", 0.9f)),
eps(OperatorBase::GetSingleArg<float>("eps", float(1e-3))), eps(OperatorBase::GetSingleArg<float>("eps", 1e-3f)),
r_max(OperatorBase::GetSingleArg<float>("r_max", float(3.0))), r_max(OperatorBase::GetSingleArg<float>("r_max", 3.f)),
d_max(OperatorBase::GetSingleArg<float>("d_max", float(5.0))), d_max(OperatorBase::GetSingleArg<float>("d_max", 5.f)),
t_delta(OperatorBase::GetSingleArg<float>("t_delta", float(1.0))), t_delta(OperatorBase::GetSingleArg<float>("t_delta", 1.f)),
use_stats(OperatorBase::GetSingleArg<int>("use_stats", -1)), 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")) { mode(OperatorBase::GetSingleArg<string>("mode", "DEFAULT")) {
if (axis != -1) if (axis != -1)
CHECK_EQ(axis, 1) CHECK_EQ(axis, 1)
......
Dragon/include/operators/norm/group_norm_op.h
...@@ -23,7 +23,7 @@ class GroupNormOp : public Operator<Context> { ...@@ -23,7 +23,7 @@ class GroupNormOp : public Operator<Context> {
: Operator<Context>(op_def, ws), : Operator<Context>(op_def, ws),
group(OperatorBase::GetSingleArg<int>("group", 32)), group(OperatorBase::GetSingleArg<int>("group", 32)),
axis(OperatorBase::GetSingleArg<int>("axis", -1)), axis(OperatorBase::GetSingleArg<int>("axis", -1)),
eps(OperatorBase::GetSingleArg<float>("eps", float(1e-3))) { eps(OperatorBase::GetSingleArg<float>("eps", 1e-3f)) {
if (axis != -1) if (axis != -1)
CHECK_EQ(axis, 1) CHECK_EQ(axis, 1)
<< "\nThe axis can only be set to 1."; << "\nThe axis can only be set to 1.";
...@@ -77,7 +77,7 @@ class FusedGroupNormOp : public Operator<Context> { ...@@ -77,7 +77,7 @@ class FusedGroupNormOp : public Operator<Context> {
: Operator<Context>(op_def, ws), : Operator<Context>(op_def, ws),
group(OperatorBase::GetSingleArg<int>("group", 32)), group(OperatorBase::GetSingleArg<int>("group", 32)),
axis(OperatorBase::GetSingleArg<int>("axis", -1)), 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); USE_OPERATOR_FUNCTIONS(Context);
void Setup(); void Setup();
...@@ -100,8 +100,7 @@ class FusedGroupNormGradientOp : public Operator<Context> { ...@@ -100,8 +100,7 @@ class FusedGroupNormGradientOp : public Operator<Context> {
FusedGroupNormGradientOp(const OperatorDef& op_def, Workspace* ws) FusedGroupNormGradientOp(const OperatorDef& op_def, Workspace* ws)
: Operator<Context>(op_def, ws), : Operator<Context>(op_def, ws),
group(OperatorBase::GetSingleArg<int>("group", 32)), group(OperatorBase::GetSingleArg<int>("group", 32)),
axis(OperatorBase::GetSingleArg<int>("axis", -1)), axis(OperatorBase::GetSingleArg<int>("axis", -1)) {}
eps(OperatorBase::GetSingleArg<float>("eps", float(1e-3))) {}
USE_OPERATOR_FUNCTIONS(Context); USE_OPERATOR_FUNCTIONS(Context);
void Setup(); void Setup();
...@@ -110,7 +109,6 @@ class FusedGroupNormGradientOp : public Operator<Context> { ...@@ -110,7 +109,6 @@ class FusedGroupNormGradientOp : public Operator<Context> {
template <typename T> void RunWithType(); template <typename T> void RunWithType();
protected: protected:
float eps;
Tensor num_by_chans; Tensor num_by_chans;
Tensor* multiplier, *num_multiplier, *spatial_multiplier, *cgs_multiplier; Tensor* multiplier, *num_multiplier, *spatial_multiplier, *cgs_multiplier;
Tensor* mean, *var, *stddev, *x_norm; Tensor* mean, *var, *stddev, *x_norm;
......
Dragon/include/operators/norm/instance_norm_op.h
...@@ -22,7 +22,7 @@ class InstanceNormOp : public Operator<Context> { ...@@ -22,7 +22,7 @@ class InstanceNormOp : public Operator<Context> {
InstanceNormOp(const OperatorDef& op_def, Workspace* ws) InstanceNormOp(const OperatorDef& op_def, Workspace* ws)
: Operator<Context>(op_def, ws), : Operator<Context>(op_def, ws),
axis(OperatorBase::GetSingleArg<int>("axis", -1)), axis(OperatorBase::GetSingleArg<int>("axis", -1)),
eps(OperatorBase::GetSingleArg<float>("eps", float(1e-3))) { eps(OperatorBase::GetSingleArg<float>("eps", 1e-3f)) {
if (axis != -1) if (axis != -1)
CHECK_EQ(axis, 1) CHECK_EQ(axis, 1)
<< "\nThe axis can only be set to 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> { ...@@ -23,7 +23,7 @@ class L2NormOp final : public Operator<Context> {
: Operator<Context>(op_def, ws), : Operator<Context>(op_def, ws),
axis(OperatorBase::GetSingleArg<int>("axis", 0)), axis(OperatorBase::GetSingleArg<int>("axis", 0)),
num_axes(OperatorBase::GetSingleArg<int>("num_axes", -1)), 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")) {} mode(OperatorBase::GetSingleArg<string>("mode", "SUM")) {}
USE_OPERATOR_FUNCTIONS(Context); USE_OPERATOR_FUNCTIONS(Context);
......
Dragon/include/operators/update/sgd_update_op.h
...@@ -20,14 +20,15 @@ template <class Context> ...@@ -20,14 +20,15 @@ template <class Context>
class SGDUpdateOp final : public UpdateOpBase<Context> { class SGDUpdateOp final : public UpdateOpBase<Context> {
public: public:
SGDUpdateOp(const OperatorDef& op_def, Workspace* ws) 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_OPERATOR_FUNCTIONS(Context);
USE_UPDATER_FUNCTIONS(Context); USE_UPDATER_FUNCTIONS(Context);
void ComputeRunWithFloat() override; void ComputeRunWithFloat() override;
protected: protected:
float lr, momentum; float old_lr, lr, momentum, correction;
}; };
} // namespace dragon } // namespace dragon
......
Dragon/include/utils/op_kernel.h
...@@ -286,6 +286,11 @@ void SparseSoftmaxFocalLossGrad(const int count, ...@@ -286,6 +286,11 @@ void SparseSoftmaxFocalLossGrad(const int count,
Tensor* ignore, Tensor* ignore,
T* dx); T* dx);
/******************** misc.dtype ********************/
template <typename Ta, typename Tb, class Context>
void TypeA2B(const int count, const Ta* a, Tb* b);
/******************** misc.image_data ********************/ /******************** misc.image_data ********************/
template <typename Tx, typename Ty, class Context> template <typename Tx, typename Ty, class Context>
...@@ -308,7 +313,7 @@ void Arange(const int count, ...@@ -308,7 +313,7 @@ void Arange(const int count,
const int step, const int step,
T* y); T* y);
/******************** ndarray.argmax ********************/ /******************** ndarray.argreduce ********************/
template <typename T, class Context> template <typename T, class Context>
void Argmax(const int count, void Argmax(const int count,
...@@ -316,9 +321,8 @@ void Argmax(const int count, ...@@ -316,9 +321,8 @@ void Argmax(const int count,
const int inner_dim, const int inner_dim,
const int top_k, const int top_k,
const T* x, const T* x,
T* y); int64_t* indices,
T* values);
/******************** ndarray.argmin ********************/
template <typename T, class Context> template <typename T, class Context>
void Argmin(const int count, void Argmin(const int count,
...@@ -326,7 +330,8 @@ void Argmin(const int count, ...@@ -326,7 +330,8 @@ void Argmin(const int count,
const int inner_dim, const int inner_dim,
const int top_k, const int top_k,
const T* x, const T* x,
T* y); int64_t* indices,
T* values);
/******************** ndarray.gather ********************/ /******************** ndarray.gather ********************/
......
Dragon/python/dragon/core/gradient_maker.py
...@@ -181,7 +181,7 @@ class GraphGradientMaker(object): ...@@ -181,7 +181,7 @@ class GraphGradientMaker(object):
# Append ops # Append ops
if not is_skip: if not is_skip:
# + Gen Op # --> GenOp
if len(gen_grads) > 0: if len(gen_grads) > 0:
op_inputs = []; op_outputs = []; values = [] op_inputs = []; op_outputs = []; values = []
for item in gen_grads: for item in gen_grads:
...@@ -193,7 +193,7 @@ class GraphGradientMaker(object): ...@@ -193,7 +193,7 @@ class GraphGradientMaker(object):
if forward_op.HasField('device_option'): if forward_op.HasField('device_option'):
gen_op.device_option.CopyFrom(forward_op.device_option) gen_op.device_option.CopyFrom(forward_op.device_option)
backward_ops.append(gen_op) backward_ops.append(gen_op)
# + Grad Op # --> GradOp
for g_op in g_ops: for g_op in g_ops:
g_op.name = _op_name(None if auto_names else 'runtime') g_op.name = _op_name(None if auto_names else 'runtime')
backward_ops.append(g_op) backward_ops.append(g_op)
...@@ -204,8 +204,10 @@ class GraphGradientMaker(object): ...@@ -204,8 +204,10 @@ class GraphGradientMaker(object):
original_idx = -1 original_idx = -1
for g_input_idx, g_input in enumerate(g_inputs): for g_input_idx, g_input in enumerate(g_inputs):
if g_output == g_input: original_idx = g_input_idx 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 original_idx == -1: continue
if g_output in g_op.input: continue
# Found a split branch
original_name = forward_op.input[original_idx] original_name = forward_op.input[original_idx]
if inputs_count[original_name] > 1: if inputs_count[original_name] > 1:
# Split # Split
......
Dragon/python/dragon/core/tensor.py
...@@ -340,6 +340,34 @@ class Tensor(object): ...@@ -340,6 +340,34 @@ class Tensor(object):
def dtype(self, value): def dtype(self, value):
self._dtype = 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 @property
def extra_targets(self): def extra_targets(self):
"""Return or Set the extra solving targets. """Return or Set the extra solving targets.
...@@ -646,6 +674,8 @@ class Tensor(object): ...@@ -646,6 +674,8 @@ class Tensor(object):
output.shape = self.shape[:] output.shape = self.shape[:]
return output return output
__truediv__ = __div__
def __rdiv__(self, other): def __rdiv__(self, other):
"""Calculate y / x. """Calculate y / x.
...@@ -672,6 +702,8 @@ class Tensor(object): ...@@ -672,6 +702,8 @@ class Tensor(object):
output.shape = self.shape[:] output.shape = self.shape[:]
return output return output
__rtruediv__ = __rdiv__
def __neg__(self): def __neg__(self):
"""Calculate -x. """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 ...@@ -129,7 +129,9 @@ List Brief
`Reduce`_ The general reduce operator. `Reduce`_ The general reduce operator.
`Sum`_ Compute the sum along the given axis. `Sum`_ Compute the sum along the given axis.
`Mean`_ Compute the mean 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. `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. `Argmin`_ Compute the indices of minimum elements along the given axis.
`Slice`_ Slice interface of NDArray. `Slice`_ Slice interface of NDArray.
`Stack`_ Stack the inputs along the given axis. `Stack`_ Stack the inputs along the given axis.
...@@ -161,6 +163,7 @@ Misc ...@@ -161,6 +163,7 @@ Misc
================= ====================================================================== ================= ======================================================================
List Brief List Brief
================= ====================================================================== ================= ======================================================================
`AsType`_ Cast the data type of inputs to a specific one.
`Run`_ Run a custom operator. (Without GradientFlow) `Run`_ Run a custom operator. (Without GradientFlow)
`Template`_ Run a custom operator. (With GradientFlow) `Template`_ Run a custom operator. (With GradientFlow)
`Accuracy`_ Calculate the Top-K accuracy. `Accuracy`_ Calculate the Top-K accuracy.
...@@ -176,14 +179,6 @@ List Brief ...@@ -176,14 +179,6 @@ List Brief
`Proposal`_ Generate Regional Proposals, introduced by `[Ren et.al, 2015] <https://arxiv.org/abs/1506.01497>`_. `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 MPI
--- ---
================= ====================================================================== ================= ======================================================================
...@@ -266,7 +261,9 @@ List Brief ...@@ -266,7 +261,9 @@ List Brief
.. _Reduce: operators/ndarray.html#dragon.operators.ndarray.Reduce .. _Reduce: operators/ndarray.html#dragon.operators.ndarray.Reduce
.. _Sum: operators/ndarray.html#dragon.operators.ndarray.Sum .. _Sum: operators/ndarray.html#dragon.operators.ndarray.Sum
.. _Mean: operators/ndarray.html#dragon.operators.ndarray.Mean .. _Mean: operators/ndarray.html#dragon.operators.ndarray.Mean
.. _Max: operators/ndarray.html#dragon.operators.ndarray.Max
.. _Argmax: operators/ndarray.html#dragon.operators.ndarray.Argmax .. _Argmax: operators/ndarray.html#dragon.operators.ndarray.Argmax
.. _Min: operators/ndarray.html#dragon.operators.ndarray.Min
.. _Argmin: operators/ndarray.html#dragon.operators.ndarray.Argmin .. _Argmin: operators/ndarray.html#dragon.operators.ndarray.Argmin
.. _Slice: operators/ndarray.html#dragon.operators.ndarray.Slice .. _Slice: operators/ndarray.html#dragon.operators.ndarray.Slice
.. _Stack: operators/ndarray.html#dragon.operators.ndarray.Stack .. _Stack: operators/ndarray.html#dragon.operators.ndarray.Stack
...@@ -285,6 +282,7 @@ List Brief ...@@ -285,6 +282,7 @@ List Brief
.. _Copy: operators/control_flow.html#dragon.operators.control_flow.Copy .. _Copy: operators/control_flow.html#dragon.operators.control_flow.Copy
.. _Equal: operators/control_flow.html#dragon.operators.control_flow.Equal .. _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 .. _Run: operators/misc.html#dragon.operators.misc.Run
.. _Template: operators/misc.html#dragon.operators.misc.Template .. _Template: operators/misc.html#dragon.operators.misc.Template
.. _Accuracy: operators/misc.html#dragon.operators.misc.Accuracy .. _Accuracy: operators/misc.html#dragon.operators.misc.Accuracy
...@@ -293,7 +291,5 @@ List Brief ...@@ -293,7 +291,5 @@ List Brief
.. _Proposal: operators/contrib/rcnn.html#dragon.operators.contrib.rcnn.ops.Proposal .. _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 .. _MPIBroadcast: operators/mpi.html#dragon.operators.mpi.MPIBroadcast
.. _MPIGather: operators/mpi.html#dragon.operators.mpi.MPIGather .. _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): ...@@ -62,7 +62,7 @@ class BlobFetcher(Process):
# fill blobs # fill blobs
im, labels = self.Q_in.get() im, labels = self.Q_in.get()
im_blob = np.zeros(shape=([self._batch_size] + list(im.shape)), dtype=np.uint8) im_blob = 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): for ix in range(0, self._batch_size):
im_blob[ix, :, :, :], label_blob[ix, :] = im, labels im_blob[ix, :, :, :], label_blob[ix, :] = im, labels
if ix != self._batch_size - 1: im, labels = self.Q_in.get() if ix != self._batch_size - 1: im, labels = self.Q_in.get()
......
Dragon/python/dragon/io/data_batch.py
...@@ -38,10 +38,10 @@ class DataBatch(object): ...@@ -38,10 +38,10 @@ class DataBatch(object):
---------- ----------
source : str source : str
The path of database. The path of database.
multiple_nodes: boolean
Whether to split data for multiple parallel nodes.
shuffle : boolean shuffle : boolean
Whether to shuffle the data. Whether to shuffle the data.
node_step: boolean
Whether to split data for multiple parallel nodes.
num_chunks : int num_chunks : int
The number of chunks to split. Default is ``2048``. The number of chunks to split. Default is ``2048``.
chunk_size : int chunk_size : int
...@@ -105,6 +105,10 @@ class DataBatch(object): ...@@ -105,6 +105,10 @@ class DataBatch(object):
if kwargs.get('max_random_scale', 1.0) - \ if kwargs.get('max_random_scale', 1.0) - \
kwargs.get('min_random_scale', 1.0) != 0: kwargs.get('min_random_scale', 1.0) != 0:
self._num_transformers += 1 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._num_transformers = min(self._num_transformers, self._max_transformers)
self._batch_size = kwargs.get('batch_size', 100) self._batch_size = kwargs.get('batch_size', 100)
...@@ -127,8 +131,8 @@ class DataBatch(object): ...@@ -127,8 +131,8 @@ class DataBatch(object):
num_parts = self._num_readers num_parts = self._num_readers
part_idx = i part_idx = i
if self._readers[i]._use_shuffle \ if self._readers[i]._multiple_nodes or \
or self._readers[i]._use_step: self._readers[i]._use_shuffle:
num_parts *= group_size num_parts *= group_size
part_idx += local_rank * self._num_readers part_idx += local_rank * self._num_readers
......
Dragon/python/dragon/io/data_reader.py
...@@ -14,6 +14,7 @@ from __future__ import division ...@@ -14,6 +14,7 @@ from __future__ import division
from __future__ import print_function from __future__ import print_function
import math import math
import numpy as np
import numpy.random as npr import numpy.random as npr
from multiprocessing import Process from multiprocessing import Process
...@@ -34,10 +35,12 @@ class DataReader(Process): ...@@ -34,10 +35,12 @@ class DataReader(Process):
---------- ----------
source : str source : str
The path of database. The path of database.
multiple_nodes: boolean
Whether to split data for multiple parallel nodes.
shuffle : boolean shuffle : boolean
Whether to shuffle the data. Whether to shuffle the data.
node_step: boolean instance_chunk : boolean
Whether to split data for multiple parallel nodes. Whether to limit each chunk with at most 1 instance.
num_chunks : int num_chunks : int
The number of chunks to split. Default is ``2048``. The number of chunks to split. Default is ``2048``.
chunk_size : int chunk_size : int
...@@ -46,8 +49,9 @@ class DataReader(Process): ...@@ -46,8 +49,9 @@ class DataReader(Process):
""" """
super(DataReader, self).__init__() super(DataReader, self).__init__()
self._source = kwargs.get('source', '') self._source = kwargs.get('source', '')
self._multiple_nodes = kwargs.get('multiple_nodes', False)
self._use_shuffle = kwargs.get('shuffle', 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._num_chunks = kwargs.get('num_chunks', 2048)
self._chunk_size = kwargs.get('chunk_size', -1) self._chunk_size = kwargs.get('chunk_size', -1)
...@@ -106,9 +110,8 @@ class DataReader(Process): ...@@ -106,9 +110,8 @@ class DataReader(Process):
None None
""" """
if self._use_shuffle or self._use_step: if self._multiple_nodes or self._use_shuffle:
if self._use_shuffle: if self._use_shuffle: self._perm = npr.permutation(self._num_shuffle_parts)
self._perm = npr.permutation(self._num_shuffle_parts)
self._cur_chunk_idx = 0 self._cur_chunk_idx = 0
self._start_idx = int(self._part_idx * self._num_shuffle_parts + self._perm[self._cur_chunk_idx]) self._start_idx = int(self._part_idx * self._num_shuffle_parts + self._perm[self._cur_chunk_idx])
self._start_idx = int(self._start_idx * self._chunk_size) self._start_idx = int(self._start_idx * self._chunk_size)
...@@ -168,6 +171,11 @@ class DataReader(Process): ...@@ -168,6 +171,11 @@ class DataReader(Process):
self._db_size = int(self._db.get('size')) self._db_size = int(self._db.get('size'))
self._db_zfill = int(self._db.get('zfill')) self._db_zfill = int(self._db.get('zfill'))
self._epoch_size = int(self._db_size / self._num_parts + 1) 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 # search a optimal chunk size by chunks
if self._chunk_size == -1: if self._chunk_size == -1:
max_chunk_size = self._db._total_size / ((self._num_chunks * (1 << 20))) max_chunk_size = self._db._total_size / ((self._num_chunks * (1 << 20)))
...@@ -177,7 +185,7 @@ class DataReader(Process): ...@@ -177,7 +185,7 @@ class DataReader(Process):
self._num_shuffle_parts = int(math.ceil(self._db._total_size * 1.1 / self._num_shuffle_parts = int(math.ceil(self._db._total_size * 1.1 /
(self._num_parts * self._chunk_size << 20))) (self._num_parts * self._chunk_size << 20)))
self._chunk_size = int(self._db_size / self._num_shuffle_parts / self._num_parts + 1) self._chunk_size = int(self._db_size / self._num_shuffle_parts / self._num_parts + 1)
self._perm = npr.permutation(self._num_shuffle_parts) self._perm = np.arange(self._num_shuffle_parts)
# init env # init env
self.reset() self.reset()
...@@ -187,5 +195,6 @@ class DataReader(Process): ...@@ -187,5 +195,6 @@ class DataReader(Process):
self.Q_out.put(self.element()) self.Q_out.put(self.element())
self.next_record() self.next_record()
if self._cur_idx >= self._end_idx: 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() 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 ...@@ -13,6 +13,7 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import sys
import numpy as np import numpy as np
import numpy.random as npr import numpy.random as npr
from multiprocessing import Process from multiprocessing import Process
...@@ -28,7 +29,7 @@ try: ...@@ -28,7 +29,7 @@ try:
import PIL.Image import PIL.Image
import PIL.ImageEnhance import PIL.ImageEnhance
except ImportError as e: 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): class DataTransformer(Process):
...@@ -103,12 +104,20 @@ class DataTransformer(Process): ...@@ -103,12 +104,20 @@ class DataTransformer(Process):
random_scale = npr.uniform() * (self._max_random_scale - self._min_random_scale) \ random_scale = npr.uniform() * (self._max_random_scale - self._min_random_scale) \
+ self._min_random_scale + self._min_random_scale
if random_scale != 1.0: if random_scale != 1.0:
if sys.version_info >= (3, 0):
im = cv2.resize(im, None, interpolation=cv2.INTER_LINEAR, im = cv2.resize(im, None, interpolation=cv2.INTER_LINEAR,
fx=random_scale, fy=random_scale) 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 # random crop
if self._crop_size > 0: if self._crop_size > 0:
if self._phase== 'TRAIN': if self._phase == 'TRAIN':
h_off = npr.randint(im.shape[0] - self._crop_size + 1) h_off = npr.randint(im.shape[0] - self._crop_size + 1)
w_off = npr.randint(im.shape[1] - self._crop_size + 1) w_off = npr.randint(im.shape[1] - self._crop_size + 1)
else: else:
...@@ -152,7 +161,12 @@ class DataTransformer(Process): ...@@ -152,7 +161,12 @@ class DataTransformer(Process):
self._padding : self._padding + im.shape[1], :] = im self._padding : self._padding + im.shape[1], :] = im
im = pad_img 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): def run(self):
"""Start the process. """Start the process.
......
Dragon/python/dragon/operators/activation.py
...@@ -200,7 +200,9 @@ def Tanh(inputs, **kwargs): ...@@ -200,7 +200,9 @@ def Tanh(inputs, **kwargs):
def Dropout(inputs, prob=0.5, scale=True, **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 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 ...@@ -16,6 +16,49 @@ from __future__ import print_function
from . import * 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): def Run(inputs, module, op, param_str='', nout=1, **kwargs):
"""Run a custom operator. (Without GradientFlow) """Run a custom operator. (Without GradientFlow)
......
Dragon/python/dragon/operators/ndarray.py
...@@ -235,7 +235,6 @@ def Concat(inputs, axis=1, **kwargs): ...@@ -235,7 +235,6 @@ def Concat(inputs, axis=1, **kwargs):
""" """
CheckInputs(inputs, 1, INT_MAX) CheckInputs(inputs, 1, INT_MAX)
arguments = ParseArguments(locals()) arguments = ParseArguments(locals())
arguments['num_input'] = len(inputs)
output = Tensor.CreateOperator(nout=1, op_type='Concat', **arguments) output = Tensor.CreateOperator(nout=1, op_type='Concat', **arguments)
...@@ -337,6 +336,26 @@ def Mean(inputs, axis=-1, keep_dims=False, **kwargs): ...@@ -337,6 +336,26 @@ def Mean(inputs, axis=-1, keep_dims=False, **kwargs):
return Reduce(inputs, axis, 'MEAN', keep_dims, **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): def Argmax(inputs, axis=-1, top_k=1, keep_dims=False, **kwargs):
"""Compute the indices of maximum elements along the given axis. """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): ...@@ -357,17 +376,30 @@ def Argmax(inputs, axis=-1, top_k=1, keep_dims=False, **kwargs):
The indices. The indices.
""" """
CheckInputs(inputs, 1) return _ArgReduce(inputs, axis, 'ARGMAX', top_k, keep_dims, **kwargs)
arguments = ParseArguments(locals())
output = Tensor.CreateOperator(nout=1, op_type='Argmax', **arguments)
if inputs.shape is not None: def Max(inputs, axis=-1, top_k=1, keep_dims=False, **kwargs):
output.shape = inputs.shape[:] """Compute the values of maximum elements along the given axis.
if top_k > 1: output.shape[axis] = top_k
else: del output.shape[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): 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): ...@@ -390,17 +422,30 @@ def Argmin(inputs, axis=-1, top_k=1, keep_dims=False, **kwargs):
The indices. The indices.
""" """
CheckInputs(inputs, 1) return _ArgReduce(inputs, axis, 'ARGMIN', top_k, keep_dims, **kwargs)
arguments = ParseArguments(locals())
output = Tensor.CreateOperator(nout=1, op_type='Argmin', **arguments)
if inputs.shape is not None: def Min(inputs, axis=-1, top_k=1, keep_dims=False, **kwargs):
output.shape = inputs.shape[:] """Compute the values of minimum elements along the given axis.
if top_k > 1: output.shape[axis] = top_k
else: del output.shape[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): def Transpose(inputs, perms=None, **kwargs):
......
Dragon/python/dragon/ops.py
...@@ -21,7 +21,6 @@ from .operators import activation as act ...@@ -21,7 +21,6 @@ from .operators import activation as act
from .operators import arithmetic as math from .operators import arithmetic as math
from .operators import control_flow from .operators import control_flow
from .operators import misc as misc from .operators import misc as misc
from .operators import cast
from .operators import mpi from .operators import mpi
from .operators import ndarray from .operators import ndarray
from .operators import norm from .operators import norm
...@@ -110,7 +109,9 @@ Crop = ndarray.Crop ...@@ -110,7 +109,9 @@ Crop = ndarray.Crop
Reduce = ndarray.Reduce Reduce = ndarray.Reduce
Sum = ndarray.Sum Sum = ndarray.Sum
Mean = ndarray.Mean Mean = ndarray.Mean
Max = ndarray.Max
Argmax = ndarray.Argmax Argmax = ndarray.Argmax
Min = ndarray.Min
Argmin = ndarray.Argmin Argmin = ndarray.Argmin
Slice = ndarray.Slice Slice = ndarray.Slice
Stack = ndarray.Stack Stack = ndarray.Stack
...@@ -131,15 +132,13 @@ Copy = control_flow.Copy ...@@ -131,15 +132,13 @@ Copy = control_flow.Copy
Equal = control_flow.Equal Equal = control_flow.Equal
# misc # misc
AsType = misc.AsType
Run = misc.Run Run = misc.Run
Template = misc.Template Template = misc.Template
Accuracy = misc.Accuracy Accuracy = misc.Accuracy
StopGradient = misc.StopGradient StopGradient = misc.StopGradient
MovingAverage = misc.MovingAverage MovingAverage = misc.MovingAverage
# cast
FloatToHalf = cast.FloatToHalf
# mpi # mpi
MPIBroadcast = mpi.MPIBroadcast MPIBroadcast = mpi.MPIBroadcast
MPIGather = mpi.MPIGather MPIGather = mpi.MPIGather
......
Dragon/python/dragon/vm/caffe/proto/caffe.proto
...@@ -38,6 +38,7 @@ message Datum { ...@@ -38,6 +38,7 @@ message Datum {
repeated float float_data = 6; repeated float float_data = 6;
// If true data contains an encoded image that need to be decoded // If true data contains an encoded image that need to be decoded
optional bool encoded = 7 [default = false]; optional bool encoded = 7 [default = false];
repeated int32 labels = 8;
} }
message FillerParameter { message FillerParameter {
...@@ -1459,8 +1460,8 @@ message NormalizeParameter { ...@@ -1459,8 +1460,8 @@ message NormalizeParameter {
} }
message ParallelParameter { message ParallelParameter {
optional bool shuffle = 1 [default = false]; optional bool multiple_nodes = 1 [default = false];
optional bool node_step = 2 [default = false]; optional bool shuffle = 2 [default = false];
optional bool partition = 3 [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 ...@@ -36,6 +36,8 @@ from dragon.config import option
from .autograd.expression import Expression from .autograd.expression import Expression
from .autograd.grad_mode import is_grad_enabled from .autograd.grad_mode import is_grad_enabled
from .constants import CTX_TO_DEVICE_OPTION 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): def RunOperator(inputs, outputs, meta, auto_grad=True, **kwargs):
...@@ -44,36 +46,41 @@ 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: if len(outputs) == 0:
raise ValueError('The num of outputs should be at least 1.') 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] for ix, output in enumerate(outputs):
outputs_name = [output.name for output in 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 engine_type = meta[0]; persistent_key = None
if engine_type == 'ONCE': if engine_type == 'ONCE':
# OpType + CTX -> Op # ++ OpType + CTX -> Op
op_type, ctx = meta[1:] op_type, ctx = meta[1:]
if ctx is None: raise ValueError('Excepted a context, got None.') if ctx is None: raise ValueError('Excepted a context, got None.')
op = pb.MakeOperatorDef(op_type, inputs_name, outputs_name, name='runtime', 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': elif engine_type == 'PERSISTENT':
# Key + Inputs + Outputs -> Op # ++ Key + Inputs + Outputs -> Op
persistent_key, meta_op = meta[1:] persistent_key, meta_op = meta[1:]
op = pb.MutableOperatorDef(meta_op, inputs_name, outputs_name) op = pb.MutableOperatorDef(meta_op, inputs_name, outputs_name)
else: else:
raise ValueError('Unknown executing engine: {}.'.format(engine_type)) raise ValueError('Unknown executing engine: {}.'.format(engine_type))
# ~ Auto-Grad ~ # + Auto-Grad
requires_grad = False
for input in inputs:
if input.requires_grad:
requires_grad = True
if len(inputs) > 0 and auto_grad: if len(inputs) > 0 and auto_grad:
input_expressions = [] input_expressions = []
if requires_grad and is_grad_enabled(): if requires_grad:
ignored_grads = set() ignored_grads = set()
# + Trace outputs # ++ Trace outputs
for input in inputs: for input in inputs:
input_expressions.append(input._expr) input_expressions.append(input._expr)
if input._ignored_grads: if input._ignored_grads:
...@@ -88,11 +95,11 @@ def RunOperator(inputs, outputs, meta, auto_grad=True, **kwargs): ...@@ -88,11 +95,11 @@ def RunOperator(inputs, outputs, meta, auto_grad=True, **kwargs):
if len(ignored_grads) > 0: if len(ignored_grads) > 0:
outputs[ix]._ignored_grads = ignored_grads outputs[ix]._ignored_grads = ignored_grads
else: else:
# + Reset status # ++ Reset status
for ix in range(len(outputs)): for ix in range(len(outputs)):
outputs[ix]._requires_grad = False outputs[ix]._requires_grad = False
# Run # + Run
if option['log_optimized_graph'] or option['log_meta_graph']: if option['log_optimized_graph'] or option['log_meta_graph']:
from dragon.config import logger from dragon.config import logger
logger.info('>>>>>>>>>>>>>>>>>> Forward Flow <<<<<<<<<<<<<<<<<<\n') logger.info('>>>>>>>>>>>>>>>>>> Forward Flow <<<<<<<<<<<<<<<<<<\n')
...@@ -104,7 +111,7 @@ def RunOperator(inputs, outputs, meta, auto_grad=True, **kwargs): ...@@ -104,7 +111,7 @@ def RunOperator(inputs, outputs, meta, auto_grad=True, **kwargs):
dg.workspace.RunPersistentOp(persistent_key, dg.workspace.RunPersistentOp(persistent_key,
op.name, inputs_name, outputs_name) op.name, inputs_name, outputs_name)
# Returns # + Returns
if len(outputs) > 1: return outputs if len(outputs) > 1: return outputs
elif len(outputs) == 1: return outputs[0] elif len(outputs) == 1: return outputs[0]
else: return None 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 ...@@ -30,14 +30,14 @@ from dragon.config import logger
from dragon.vm.torch.environ import \ from dragon.vm.torch.environ import \
add_submodule, get_module_name 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.tensor_pool import TPool
from dragon.vm.torch.execute_engine import RunOperator from dragon.vm.torch.execute_engine import RunOperator
class Module(object): class Module(object):
def __init__(self): def __init__(self):
self._modules = {} self._modules = OrderedDict()
self._parameters = OrderedDict() self._parameters = OrderedDict()
self._buffers = OrderedDict() self._buffers = OrderedDict()
self._persistent_key = self._op = None self._persistent_key = self._op = None
...@@ -96,22 +96,24 @@ class Module(object): ...@@ -96,22 +96,24 @@ class Module(object):
else: else:
object.__setattr__(self, key, value) 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: if destination is None:
destination = OrderedDict() destination = OrderedDict()
for name, param in self._parameters.items(): for name, param in self._parameters.items():
if param is not None: 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(): for name, buf in self._buffers.items():
if buf is not None: 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(): for name, module in self._modules.items():
if module is not None: 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 return destination
def load_state_dict(self, state_dict, strict=True): 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): def submodule_key_mismatch(full_name, is_missing):
module = self module = self
names = full_name.split(".") names = full_name.split(".")
...@@ -198,6 +200,14 @@ class Module(object): ...@@ -198,6 +200,14 @@ class Module(object):
else: else:
self._parameters[name] = param 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): def __call__(self, *args, **kwargs):
with dg.name_scope(get_module_name(self)): with dg.name_scope(get_module_name(self)):
return self.forward(*args) return self.forward(*args)
...@@ -238,6 +248,23 @@ class Module(object): ...@@ -238,6 +248,23 @@ class Module(object):
memo.add(module) memo.add(module)
yield name, 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=''): def named_parameters(self, memo=None, prefix=''):
"""Returns an iterator over module parameters. """Returns an iterator over module parameters.
...@@ -315,7 +342,7 @@ class Module(object): ...@@ -315,7 +342,7 @@ class Module(object):
raise NotImplementedError() raise NotImplementedError()
def register_output(self, dtype='float32'): def register_output(self, dtype='float32'):
return RuntimeTensor(dtype=dtype, ctx=self._ctx) return (dtype, self._ctx)
def unify_devices(self, inputs): def unify_devices(self, inputs):
for ix, t in enumerate(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 ...@@ -18,8 +18,10 @@ from dragon.vm.torch.module import Module
from dragon.vm.torch.tensor import Parameter from dragon.vm.torch.tensor import Parameter
from .modules.conv import Conv2d from .modules.conv import Conv2d
from .modules.pooling import MaxPool2d, AvgPool2d from .modules.pooling import MaxPool2d, AvgPool2d
from .modules.activation import ReLU from .modules.activation import ReLU, Softmax
from .modules.linear import Linear from .modules.linear import Linear
from .modules.loss import CrossEntropyLoss from .modules.loss import CrossEntropyLoss
from .modules.container import Container, Sequential, ModuleList
from .modules.batchnorm import BatchNorm1d, BatchNorm2d, BatchNorm3d from .modules.batchnorm import BatchNorm1d, BatchNorm2d, BatchNorm3d
from .modules.dropout import Dropout, Dropout2d, Dropout3d
from . import init 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 ...@@ -19,6 +19,7 @@ from __future__ import print_function
import math import math
import warnings import warnings
from dragon.vm.torch.autograd.grad_mode import no_grad
def calculate_gain(nonlinearity, param=None): def calculate_gain(nonlinearity, param=None):
...@@ -43,14 +44,17 @@ def calculate_gain(nonlinearity, param=None): ...@@ -43,14 +44,17 @@ def calculate_gain(nonlinearity, param=None):
def uniform_(tensor, a=0, b=1): def uniform_(tensor, a=0, b=1):
with no_grad():
return tensor.uniform_(a, b) return tensor.uniform_(a, b)
def normal_(tensor, mean=0, std=1): def normal_(tensor, mean=0, std=1):
with no_grad():
return tensor.normal_(mean, std) return tensor.normal_(mean, std)
def constant_(tensor, val): def constant_(tensor, val):
with no_grad():
return tensor.fill_(val) return tensor.fill_(val)
...@@ -61,8 +65,8 @@ def dirac_(tensor): ...@@ -61,8 +65,8 @@ def dirac_(tensor):
sizes = tensor.size() sizes = tensor.size()
min_dim = min(sizes[0], sizes[1]) min_dim = min(sizes[0], sizes[1])
with no_grad():
tensor.zero_() tensor.zero_()
for d in range(min_dim): for d in range(min_dim):
if dimensions == 3: # Temporal convolution if dimensions == 3: # Temporal convolution
tensor[d, d, tensor.size(2) // 2] = 1 tensor[d, d, tensor.size(2) // 2] = 1
...@@ -86,7 +90,9 @@ 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) num_output_fmaps = tensor.size(0)
receptive_field_size = 1 receptive_field_size = 1
if tensor.dim() > 2: 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_in = num_input_fmaps * receptive_field_size
fan_out = num_output_fmaps * receptive_field_size fan_out = num_output_fmaps * receptive_field_size
...@@ -97,12 +103,14 @@ def xavier_uniform_(tensor, gain=1): ...@@ -97,12 +103,14 @@ def xavier_uniform_(tensor, gain=1):
fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor) fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
std = gain * math.sqrt(2.0 / (fan_in + fan_out)) std = gain * math.sqrt(2.0 / (fan_in + fan_out))
a = math.sqrt(3.0) * std # Calculate uniform bounds from standard deviation a = math.sqrt(3.0) * std # Calculate uniform bounds from standard deviation
with no_grad():
return tensor.uniform_(-a, a) return tensor.uniform_(-a, a)
def xavier_normal_(tensor, gain=1): def xavier_normal_(tensor, gain=1):
fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor) fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
std = gain * math.sqrt(2.0 / (fan_in + fan_out)) std = gain * math.sqrt(2.0 / (fan_in + fan_out))
with no_grad():
return tensor.normal_(0, std) return tensor.normal_(0, std)
...@@ -121,6 +129,7 @@ def kaiming_uniform_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu'): ...@@ -121,6 +129,7 @@ def kaiming_uniform_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu'):
gain = calculate_gain(nonlinearity, a) gain = calculate_gain(nonlinearity, a)
std = gain / math.sqrt(fan) std = gain / math.sqrt(fan)
bound = math.sqrt(3.0) * std # Calculate uniform bounds from standard deviation bound = math.sqrt(3.0) * std # Calculate uniform bounds from standard deviation
with no_grad():
return tensor.uniform_(-bound, bound) return tensor.uniform_(-bound, bound)
...@@ -128,6 +137,7 @@ def kaiming_normal_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu'): ...@@ -128,6 +137,7 @@ def kaiming_normal_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu'):
fan = _calculate_correct_fan(tensor, mode) fan = _calculate_correct_fan(tensor, mode)
gain = calculate_gain(nonlinearity, a) gain = calculate_gain(nonlinearity, a)
std = gain / math.sqrt(fan) std = gain / math.sqrt(fan)
with no_grad():
return tensor.normal_(0, std) return tensor.normal_(0, std)
......
Dragon/python/dragon/vm/torch/nn/modules/activation.py
...@@ -33,3 +33,26 @@ class ReLU(Module): ...@@ -33,3 +33,26 @@ class ReLU(Module):
inputs = [x]; self.unify_devices(inputs) inputs = [x]; self.unify_devices(inputs)
outputs = [x if self._inplace else self.register_output(x.dtype)] outputs = [x if self._inplace else self.register_output(x.dtype)]
return self.run(inputs, outputs) 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 ...@@ -20,7 +20,7 @@ from dragon.vm.torch.module import RunOperator
class _BatchNorm(Module): 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): track_running_stats=True):
super(_BatchNorm, self).__init__() super(_BatchNorm, self).__init__()
self.num_features = num_features 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): ...@@ -102,7 +102,7 @@ class Conv2d(_ConvNd):
False, _pair(0), groups, bias) False, _pair(0), groups, bias)
def forward(self, input): 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) self.unify_devices(inputs)
outputs = [self.register_output(input.dtype)] outputs = [self.register_output(input.dtype)]
return self.run(inputs, outputs) 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): ...@@ -49,7 +49,7 @@ class Linear(Module):
self.bias.data.uniform_(-stdv, stdv) self.bias.data.uniform_(-stdv, stdv)
def forward(self, input): 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) self.unify_devices(inputs)
outputs = [self.register_output(input.dtype)] outputs = [self.register_output(input.dtype)]
return self.run(inputs, outputs) 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 ( ...@@ -19,5 +19,5 @@ from .arithmetic import (
) )
from .ndarray 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 ...@@ -13,39 +13,37 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import numpy as np
import dragon as dg
from dragon.vm.torch.tensor import Tensor from dragon.vm.torch.tensor import Tensor
from dragon.vm.torch.tensor_uitls import from_dragon from dragon.vm.torch.ops.primitive import MakeContext, WrapScalar
from dragon.vm.torch.ops.primitive import MakeContext
from dragon.vm.torch.ops.factory import get_module from dragon.vm.torch.ops.factory import get_module
from dragon.vm.torch.ops.modules.arithmetic import Fundamental from dragon.vm.torch.ops.modules.arithmetic import Fundamental
def _wrap_scalar(scalar, dtype): def _fundamental(input, value, op='Add', out=None):
# TODO(PhyscalX): We use (dtype/value) to hash different scalars. if not isinstance(value, Tensor):
# TODO(PhyscalX): In Dragon, set a Tensor with same dtype and shape will not deconstruct it. if not isinstance(value, (int, float)):
value = np.array([scalar], dtype=dtype) raise TypeError('Type of value should be numerical, got {}.'
t = dg.Tensor('/share/scalar/{}/{}'.format( .format(type(value)))
dtype, str(float(scalar)))).Variable() value = WrapScalar(value, input._dtype, input._ctx)
t.set_value(value)
return t.name
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, Tensor):
if not isinstance(value, (int, float)): if not isinstance(value, (int, float)):
raise TypeError('Type of value should be numerical, got {}.' raise TypeError('Type of value should be numerical, got {}.'
.format(type(value))) .format(type(value)))
# TODO(PhyscalX): We simply use a global tensor to share all scalars value = WrapScalar(value, input._dtype, input._ctx)
scalar_name = _wrap_scalar(value, dtype=input._dtype)
value = Tensor(dg_tensor=scalar_name, ctx=input._ctx, own_storage=False)
ctx = MakeContext(inputs=[input, value]) 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) 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): def add(input, value, out=None):
......
Dragon/python/dragon/vm/torch/ops/builtin.py
...@@ -16,10 +16,14 @@ from __future__ import print_function ...@@ -16,10 +16,14 @@ from __future__ import print_function
from dragon.vm.torch.tensor import Tensor from dragon.vm.torch.tensor import Tensor
from dragon.vm.torch.execute_engine import RunOperator 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.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.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): ...@@ -43,7 +47,7 @@ def copy_(self, src):
The ``self`` tensor. The ``self`` tensor.
""" """
return _copy(src, out=self) return _copy(self, src)
Tensor.copy_ = copy_ Tensor.copy_ = copy_
...@@ -69,7 +73,7 @@ def fill_(self, value): ...@@ -69,7 +73,7 @@ def fill_(self, value):
The self. The self.
""" """
return _fill(self, self.shape, value, self) return _fill(self, self.shape, value)
def uniform_(self, low=0, high=1): def uniform_(self, low=0, high=1):
...@@ -164,6 +168,10 @@ def add_(self, value): ...@@ -164,6 +168,10 @@ def add_(self, value):
return _fundamental(self, value, out=self, op='Add') return _fundamental(self, value, out=self, op='Add')
def radd(self, value):
return _rfundamental(self, value, op='RAdd')
def sub(self, value): def sub(self, value):
"""Subtract the ``self`` and ``value`` into the output tensor. """Subtract the ``self`` and ``value`` into the output tensor.
...@@ -198,6 +206,10 @@ def sub_(self, value): ...@@ -198,6 +206,10 @@ def sub_(self, value):
return _fundamental(self, value, out=self, op='Sub') return _fundamental(self, value, out=self, op='Sub')
def rsub(self, value):
return _rfundamental(self, value, op='RSub')
def mul(self, value): def mul(self, value):
"""Multiply the ``self`` and ``value`` into the output tensor. """Multiply the ``self`` and ``value`` into the output tensor.
...@@ -232,6 +244,10 @@ def mul_(self, value): ...@@ -232,6 +244,10 @@ def mul_(self, value):
return _fundamental(self, value, out=self, op='Mul') return _fundamental(self, value, out=self, op='Mul')
def rmul(self, value):
return _rfundamental(self, value, op='RMul')
def div(self, value): def div(self, value):
"""Divide the ``self`` and ``value`` into the output tensor. """Divide the ``self`` and ``value`` into the output tensor.
...@@ -266,14 +282,23 @@ def div_(self, value): ...@@ -266,14 +282,23 @@ def div_(self, value):
return _fundamental(self, value, out=self, op='Div') 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.add_ = add_ Tensor.add_ = add_
Tensor.__radd__ = radd
Tensor.sub = sub Tensor.sub = sub
Tensor.sub_ = sub_ Tensor.sub_ = sub_
Tensor.__rsub__ = rsub
Tensor.mul = mul Tensor.mul = mul
Tensor.mul_ = mul_ Tensor.mul_ = mul_
Tensor.__rmul__ = rmul
Tensor.div = div Tensor.div = div
Tensor.div_ = div_ Tensor.div_ = div_
Tensor.__rdiv__ = rdiv
Tensor.__rtruediv__ = rdiv
############################################## ##############################################
...@@ -302,15 +327,145 @@ def crop(self, starts, ends): ...@@ -302,15 +327,145 @@ def crop(self, starts, ends):
def mean(self, dim=None, keepdim=False): 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): 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 = view
Tensor.view_as = view_as Tensor.view_as = view_as
Tensor.mean = mean Tensor.mean = mean
Tensor.sum = sum Tensor.sum = sum
Tensor.max = max
Tensor.min = min
Tensor._crop = crop 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 ...@@ -14,8 +14,9 @@ from dragon.vm.torch.ops.factory import get_module
from dragon.vm.torch.ops.modules.control_flow import Copy from dragon.vm.torch.ops.modules.control_flow import Copy
def _copy(input, out=None): def _copy(dst, src):
ctx = MakeContext(inputs=[input]) if id(dst) == id(src): return dst
key = 'torch/ops/copy/{}:{}'.format(ctx[0].lower(), ctx[1], ) ctx = MakeContext(inputs=[dst])
key = 'torch/ops/copy/{}:{}'.format(ctx[0].lower(), ctx[1])
module = get_module(Copy, key, ctx) module = get_module(Copy, key, ctx)
return module.forward(input, out) return module.forward(dst, src)
\ No newline at end of file \ 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): ...@@ -33,7 +33,6 @@ class Copy(BaseModule):
'arguments': {} 'arguments': {}
} }
def forward(self, x, y): def forward(self, dst, src):
inputs = [x]; self.unify_devices(inputs) outputs = [dst]; self.unify_devices(outputs)
outputs = [y] if y else [self.register_output(x.dtype)] return self.run([src], outputs)
return self.run(inputs, outputs) \ No newline at end of file
\ 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 ...@@ -19,8 +19,7 @@ from dragon.vm.torch.ops.modules.base import BaseModule
class Reduce(BaseModule): class Reduce(BaseModule):
def __init__(self, key, ctx, **kwargs): def __init__(self, key, ctx, **kwargs):
super(Reduce, self).__init__(key, ctx, **kwargs) super(Reduce, self).__init__(key, ctx, **kwargs)
self.op_type = kwargs.get('op_type', 'Reduce') self.operation = kwargs.get('operation', 'SUM')
self.tag = kwargs.get('tag', None)
self.axis = kwargs.get('axis', -1) self.axis = kwargs.get('axis', -1)
self.keep_dims = kwargs.get('keep_dims', True) self.keep_dims = kwargs.get('keep_dims', True)
self.register_arguments() self.register_arguments()
...@@ -37,10 +36,10 @@ class Reduce(BaseModule): ...@@ -37,10 +36,10 @@ class Reduce(BaseModule):
def register_op(self): def register_op(self):
self.op_meta = { self.op_meta = {
'op_type': self.op_type, 'op_type': 'Reduce',
'n_inputs': 1, 'n_outputs': 1, 'n_inputs': 1, 'n_outputs': 1,
'arguments': { 'arguments': {
'operation': self.tag, 'operation': self.operation,
'axis': self.axis, 'axis': self.axis,
'keep_dims': self.keep_dims 'keep_dims': self.keep_dims
} }
...@@ -50,3 +49,56 @@ class Reduce(BaseModule): ...@@ -50,3 +49,56 @@ class Reduce(BaseModule):
inputs = [x]; self.unify_devices(inputs) inputs = [x]; self.unify_devices(inputs)
outputs = [y] if y else [self.register_output(x.dtype)] outputs = [y] if y else [self.register_output(x.dtype)]
return self.run(inputs, outputs) 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): ...@@ -31,20 +31,19 @@ class Fill(BaseModule):
def register_op(self): def register_op(self):
self.op_meta = { self.op_meta = {
'op_type': 'Fill', 'op_type': 'Fill',
'n_inputs': 1, 'n_outputs': 1, 'n_inputs': 0, 'n_outputs': 1,
'arguments': { 'arguments': {
'value': float(self.value), 'value': float(self.value),
'dims_desc': [d for d in self.shape] if len(self.shape) > 0 else None, 'dims_desc': [d for d in self.shape] if len(self.shape) > 0 else None,
} }
} }
def forward(self, x, shape, y=None): def forward(self, x, shape):
inputs = [x]; self.unify_devices(inputs) outputs = [x]; self.unify_devices(outputs)
outputs = [y if y else self.register_output(x.dtype)]
if shape is not None: if shape is not None:
for ix, d in enumerate(shape): for ix, d in enumerate(shape):
self.set_argument_i(self.shape[ix], d) self.set_argument_i(self.shape[ix], d)
return self.run(inputs, outputs) return self.run([], outputs)
class Reshape(BaseModule): class Reshape(BaseModule):
......
Dragon/python/dragon/vm/torch/ops/modules/update.py
...@@ -13,6 +13,7 @@ from __future__ import absolute_import ...@@ -13,6 +13,7 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import dragon.core.mpi as mpi
from dragon.vm.torch.ops.modules.base import BaseModule from dragon.vm.torch.ops.modules.base import BaseModule
...@@ -42,4 +43,40 @@ class Update(BaseModule): ...@@ -42,4 +43,40 @@ class Update(BaseModule):
} }
def forward(self, param, grad): def forward(self, param, grad):
self.unify_devices([param, grad])
return self.run([grad], [param], auto_grad=False) 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 ...@@ -13,11 +13,12 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
from dragon.vm.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.factory import get_module
from dragon.vm.torch.ops.modules.shape import Reshape, Fill 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.crop import Crop
from dragon.vm.torch.ops.modules.concat import Concat
def reshape(input, shape, shape_like=None): def reshape(input, shape, shape_like=None):
...@@ -28,22 +29,33 @@ def reshape(input, shape, shape_like=None): ...@@ -28,22 +29,33 @@ def reshape(input, shape, shape_like=None):
return module.forward(input, shape) 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) ctx = MakeContext(inputs=[input]); len_shape = len(shape)
key = 'torch/ops/fill/{}:{}/ndims:#{}/value:{}'.format( key = 'torch/ops/fill/{}:{}/ndims:#{}/value:{}'.format(
ctx[0].lower(), ctx[1], len_shape, value) ctx[0].lower(), ctx[1], len_shape, value)
module = get_module(Fill, key, ctx, len_shape=len_shape, value=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]) ctx = MakeContext(inputs=[input])
if dim is None: dim = -1 if dim is None: dim = -1; keepdim = False
key = 'torch/ops/{}/{}:{}/dim[{}]/keep_dims:{}'.format( elif dim < 0: dim = CanonicalAxis(input, dim)
op.lower() + ':{}'.format(tag.lower()) if tag else '', key = 'torch/ops/{}/{}:{}/dim[{}]/keep_dims:{}'.format(operation.lower(),
ctx[0].lower(), ctx[1], dim, int(keepdim)) ctx[0].lower(), ctx[1], dim, int(keepdim))
module = get_module(Reduce, key, ctx, op_type=op, module = get_module(Reduce, key, ctx,
tag=tag, axis=dim, keep_dims=keepdim) 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) return module.forward(input, out)
...@@ -67,7 +79,7 @@ def mean(input, dim=None, keepdim=False, out=None): ...@@ -67,7 +79,7 @@ def mean(input, dim=None, keepdim=False, out=None):
The mean-reduced tensor. 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): def sum(input, dim=None, keepdim=False, out=None):
...@@ -87,10 +99,10 @@ def sum(input, dim=None, keepdim=False, out=None): ...@@ -87,10 +99,10 @@ def sum(input, dim=None, keepdim=False, out=None):
Returns Returns
------- -------
vm.torch.Tensor 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): def argmax(input, dim=None, keepdim=False, out=None):
...@@ -110,14 +122,60 @@ def argmax(input, dim=None, keepdim=False, out=None): ...@@ -110,14 +122,60 @@ def argmax(input, dim=None, keepdim=False, out=None):
Returns Returns
------- -------
vm.torch.Tensor 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): 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 Parameters
---------- ----------
...@@ -132,11 +190,69 @@ def argmin(input, dim=None, keepdim=False, out=None): ...@@ -132,11 +190,69 @@ def argmin(input, dim=None, keepdim=False, out=None):
Returns 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 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): def _crop(input, starts, ends):
......
Dragon/python/dragon/vm/torch/ops/primitive.py
...@@ -13,6 +13,8 @@ from __future__ import absolute_import ...@@ -13,6 +13,8 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import numpy as np
import dragon as dg
from dragon.vm.torch.tensor import * from dragon.vm.torch.tensor import *
...@@ -58,3 +60,23 @@ def MakeContext(inputs=(), outputs=(), meta=None): ...@@ -58,3 +60,23 @@ def MakeContext(inputs=(), outputs=(), meta=None):
# Case #4: [...], [...] -> Refer Outputs # Case #4: [...], [...] -> Refer Outputs
# Case #5: meta -> CPU, CUDA:? # Case #5: meta -> CPU, CUDA:?
return type, device_id 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 ...@@ -15,26 +15,20 @@ from __future__ import print_function
import dragon.core.mpi as mpi 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.primitive import MakeContext
from dragon.vm.torch.ops.factory import get_module 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 if not mpi.Is_Init(): return
idx, group = mpi.AllowParallel() if not isinstance(grads, (list, tuple)): grads = [grads]
arguments = {} ctx = MakeContext(inputs=grads)
if idx != -1: mode = mpi.GetParallelMode() + '_ALLREDUCE'
arguments['mode'] = mpi.GetParallelMode() + '_ALLREDUCE' key = 'torch/ops/collective/{}:{}/{}'.format(
arguments['comm'], arguments['group'] \ ctx[0].lower(), ctx[1], mode.lower())
= mpi.CreateGroup(root=group[0], incl=group) module = get_module(Collective, key, ctx, mode=mode)
arguments['root'] = group[0] return module.forward(grads)
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)
def _update(param, grad, op_type, slot, def _update(param, grad, op_type, slot,
......
Dragon/python/dragon/vm/torch/optim/optimizer.py
...@@ -24,7 +24,7 @@ import numpy as np ...@@ -24,7 +24,7 @@ import numpy as np
import dragon as dg import dragon as dg
from dragon.vm.torch.tensor import Tensor 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 _OPTIMIZER_GROUP_UID = 0
...@@ -95,14 +95,15 @@ class Optimizer(object): ...@@ -95,14 +95,15 @@ class Optimizer(object):
for p in group['params']: for p in group['params']:
g_name = p.name + '_grad' g_name = p.name + '_grad'
if not dg.workspace.HasTensor(g_name): continue 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) params.append(p); grads.append(g)
# Feed optimizer parameters to workspace # Feed optimizer parameters to workspace
self.feed_parameters(group) self.feed_parameters(group)
# Run a all-reduce op to accumulate grads if necessary # Run a all-reduce op to accumulate grads if necessary
_allreduce_op(params, grads) _allreduce(grads)
# Run regular update ops # Run regular update ops
for p, g in zip(params, grads): 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): ...@@ -79,22 +79,28 @@ def save(obj, f, pickle_module=pickle, pickle_protocol=DEFAULT_PROTOCOL):
lambda f: _save(obj, f, pickle_module, pickle_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: try:
return pickle.load(f) return pickle_module.load(f)
except UnicodeDecodeError: 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): def load(f, map_location=None, pickle_module=pickle):
new_fd = False new_fd = False
file = None
if isinstance(f, str) or \ if isinstance(f, str) or \
(sys.version_info[0] == 2 and isinstance(f, unicode)) or \ (sys.version_info[0] == 2 and isinstance(f, unicode)) or \
(sys.version_info[0] == 3 and isinstance(f, pathlib.Path)): (sys.version_info[0] == 3 and isinstance(f, pathlib.Path)):
new_fd = True new_fd = True
file = f
f = open(f, 'rb') f = open(f, 'rb')
try: try:
return _load(f, map_location, pickle_module) return _load(f, map_location, pickle_module, file)
finally: finally:
if new_fd: if new_fd:
f.close() f.close()
\ No newline at end of file
Dragon/python/dragon/vm/torch/tensor.py
...@@ -315,6 +315,17 @@ class Tensor(object): ...@@ -315,6 +315,17 @@ class Tensor(object):
""" """
return self.div_(other) return self.div_(other)
def __neg__(self):
"""Calculate -x.
Returns
-------
vm.torch.Tensor
The output tensor.
"""
return self.mul(-1.0)
def __repr__(self): def __repr__(self):
"""Return a format str representing the internal storage. """Return a format str representing the internal storage.
...@@ -404,6 +415,9 @@ class Tensor(object): ...@@ -404,6 +415,9 @@ class Tensor(object):
if not keep_dims: ends[-1] = -1 if not keep_dims: ends[-1] = -1
return self._crop(starts, ends) return self._crop(starts, ends)
def __setitem__(self, key, value):
print(key, value.name)
def device(self): def device(self):
return self._ctx[1] return self._ctx[1]
...@@ -489,16 +503,24 @@ class Tensor(object): ...@@ -489,16 +503,24 @@ class Tensor(object):
""" """
return self._dtype return self._dtype
def type(self): def type(self, dtype=None):
"""Return the data type of this tensor. """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 Returns
------- -------
str str or vm.torch.Tensor
The data type. 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): ...@@ -520,7 +542,7 @@ class Tensor(object):
The output tensor. The output tensor.
""" """
raise NotImplementedError('Refer torch.ops.builtins.view') raise NotImplementedError('Refer torch.ops.builtin.view')
def view_as(self, other): def view_as(self, other):
"""Returns a new tensor with the same data but a different size as the given tensor. """Returns a new tensor with the same data but a different size as the given tensor.
...@@ -536,7 +558,7 @@ class Tensor(object): ...@@ -536,7 +558,7 @@ class Tensor(object):
The output tensor. The output tensor.
""" """
raise NotImplementedError('Refer torch.ops.builtins.view_as') raise NotImplementedError('Refer torch.ops.builtin.view_as')
def copy_(self, src): def copy_(self, src):
"""Copy the elements from ``src`` into this tensor and return ``self``. """Copy the elements from ``src`` into this tensor and return ``self``.
...@@ -552,7 +574,7 @@ class Tensor(object): ...@@ -552,7 +574,7 @@ class Tensor(object):
The ``self`` tensor. The ``self`` tensor.
""" """
raise NotImplementedError('Refer torch.ops.builtins.copy_') raise NotImplementedError('Refer torch.ops.builtin.copy_')
def fill_(self, value): def fill_(self, value):
"""Fills self tensor with the specified value. """Fills self tensor with the specified value.
...@@ -567,7 +589,7 @@ class Tensor(object): ...@@ -567,7 +589,7 @@ class Tensor(object):
The self. The self.
""" """
raise NotImplementedError('Refer torch.ops.builtins.fill_') raise NotImplementedError('Refer torch.ops.builtin.fill_')
def zero_(self): def zero_(self):
"""Fills self tensor with zeros. """Fills self tensor with zeros.
...@@ -600,7 +622,7 @@ class Tensor(object): ...@@ -600,7 +622,7 @@ class Tensor(object):
The self. The self.
""" """
raise NotImplementedError('Refer torch.ops.builtins.uniform_') raise NotImplementedError('Refer torch.ops.builtin.uniform_')
def normal_(self, mean=0, std=1): def normal_(self, mean=0, std=1):
"""Fill self tensor with the specified normal distribution. """Fill self tensor with the specified normal distribution.
...@@ -618,7 +640,7 @@ class Tensor(object): ...@@ -618,7 +640,7 @@ class Tensor(object):
The self. The self.
""" """
raise NotImplementedError('Refer torch.ops.builtins.normal_') raise NotImplementedError('Refer torch.ops.builtin.normal_')
def add(self, value): def add(self, value):
"""See ``torch.add()`` """See ``torch.add()``
...@@ -634,7 +656,7 @@ class Tensor(object): ...@@ -634,7 +656,7 @@ class Tensor(object):
The output tensor. The output tensor.
""" """
raise NotImplementedError('Refer torch.ops.builtins.add') raise NotImplementedError('Refer torch.ops.builtin.add')
def add_(self, value): def add_(self, value):
"""Inplace of ``torch.add()`` """Inplace of ``torch.add()``
...@@ -650,7 +672,7 @@ class Tensor(object): ...@@ -650,7 +672,7 @@ class Tensor(object):
The self. The self.
""" """
raise NotImplementedError('Refer torch.ops.builtins.add_') raise NotImplementedError('Refer torch.ops.builtin.add_')
def sub(self, value): def sub(self, value):
"""Subtract the ``self`` and ``value`` into the output tensor. """Subtract the ``self`` and ``value`` into the output tensor.
...@@ -666,7 +688,7 @@ class Tensor(object): ...@@ -666,7 +688,7 @@ class Tensor(object):
The output tensor. The output tensor.
""" """
raise NotImplementedError('Refer torch.ops.builtins.sub') raise NotImplementedError('Refer torch.ops.builtin.sub')
def sub_(self, value): def sub_(self, value):
"""Inplace of ``Tensor.sub()`` """Inplace of ``Tensor.sub()``
...@@ -682,7 +704,7 @@ class Tensor(object): ...@@ -682,7 +704,7 @@ class Tensor(object):
The self. The self.
""" """
raise NotImplementedError('Refer torch.ops.builtins.sub_') raise NotImplementedError('Refer torch.ops.builtin.sub_')
def mul(self, value): def mul(self, value):
"""Multiply the ``self`` and ``value`` into the output tensor. """Multiply the ``self`` and ``value`` into the output tensor.
...@@ -698,7 +720,7 @@ class Tensor(object): ...@@ -698,7 +720,7 @@ class Tensor(object):
The output tensor. The output tensor.
""" """
raise NotImplementedError('Refer torch.ops.builtins.mul') raise NotImplementedError('Refer torch.ops.builtin.mul')
def mul_(self, value): def mul_(self, value):
"""Inplace of ``Tensor.mul()`` """Inplace of ``Tensor.mul()``
...@@ -714,7 +736,7 @@ class Tensor(object): ...@@ -714,7 +736,7 @@ class Tensor(object):
The self. The self.
""" """
raise NotImplementedError('Refer torch.ops.builtins.mul_') raise NotImplementedError('Refer torch.ops.builtin.mul_')
def div(self, value): def div(self, value):
"""Divide the ``self`` and ``value`` into the output tensor. """Divide the ``self`` and ``value`` into the output tensor.
...@@ -730,7 +752,7 @@ class Tensor(object): ...@@ -730,7 +752,7 @@ class Tensor(object):
The output tensor. The output tensor.
""" """
raise NotImplementedError('Refer torch.ops.builtins.div') raise NotImplementedError('Refer torch.ops.builtin.div')
def div_(self, value): def div_(self, value):
"""Inplace of ``Tensor.div()`` """Inplace of ``Tensor.div()``
...@@ -746,7 +768,7 @@ class Tensor(object): ...@@ -746,7 +768,7 @@ class Tensor(object):
The self. The self.
""" """
raise NotImplementedError('Refer torch.ops.builtins.div_') raise NotImplementedError('Refer torch.ops.builtin.div_')
def mean(self, dim=None, keepdim=False): def mean(self, dim=None, keepdim=False):
"""Returns the mean of all elements or elements along the given dim. """Returns the mean of all elements or elements along the given dim.
...@@ -764,9 +786,9 @@ class Tensor(object): ...@@ -764,9 +786,9 @@ class Tensor(object):
The mean-reduced tensor. 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. """Returns the sum of all elements or elements along the given dim.
Parameters Parameters
...@@ -779,10 +801,112 @@ class Tensor(object): ...@@ -779,10 +801,112 @@ class Tensor(object):
Returns Returns
------- -------
vm.torch.Tensor 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.builtins.sum') 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.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): ...@@ -810,7 +934,7 @@ class Tensor(object):
@property @property
def grad(self): def grad(self):
g = from_dragon(self.name + '_grad', False) g = from_dragon(self.name + '_grad', False)
g._static_shape = self.shape if g: g._static_shape = self.shape
return g return g
@property @property
...@@ -836,11 +960,19 @@ class Tensor(object): ...@@ -836,11 +960,19 @@ class Tensor(object):
############################################## ##############################################
def _type2str(self): def _type2str(self):
return {'float32': 'FloatTensor', return {
'float16': 'HalfTensor',
'float32': 'FloatTensor',
'float64': 'DoubleTensor', 'float64': 'DoubleTensor',
'int32': 'IntTensor', 'int32': 'IntTensor',
'int64': 'LongTensor', 'int64': 'LongTensor',
'uint8': 'ByteTensor'}[self._dtype] 'uint8': 'ByteTensor'
}[self._dtype]
def HalfTensor(*args, **kwargs):
kwargs['dtype'] = 'float16'
return Tensor(*args, **kwargs)
def FloatTensor(*args, **kwargs): def FloatTensor(*args, **kwargs):
...@@ -869,6 +1001,7 @@ def ByteTensor(*args, **kwargs): ...@@ -869,6 +1001,7 @@ def ByteTensor(*args, **kwargs):
_DTYPE_TO_TENSOR = { _DTYPE_TO_TENSOR = {
'float16': HalfTensor,
'float32': FloatTensor, 'float32': FloatTensor,
'float64': DoubleTensor, 'float64': DoubleTensor,
'int32': IntTensor, 'int32': IntTensor,
...@@ -887,7 +1020,7 @@ def LeafTensor(shape, dtype='float32', ctx=None, requires_grad=False): ...@@ -887,7 +1020,7 @@ def LeafTensor(shape, dtype='float32', ctx=None, requires_grad=False):
return constructor(*shape, ctx=ctx, requires_grad=requires_grad) 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. """Create a torch tensor according to dtype and ctx.
Commonly used to represent the outputs that are hard to compute shape, Commonly used to represent the outputs that are hard to compute shape,
...@@ -895,7 +1028,6 @@ def RuntimeTensor(dtype='float32', ctx=None): ...@@ -895,7 +1028,6 @@ def RuntimeTensor(dtype='float32', ctx=None):
""" """
constructor = _DTYPE_TO_TENSOR[dtype] constructor = _DTYPE_TO_TENSOR[dtype]
name = TPool.get('runtime')
dg.workspace.CreateTensor(name) dg.workspace.CreateTensor(name)
return constructor(dg_tensor=name, ctx=ctx) return constructor(dg_tensor=name, ctx=ctx)
......
Dragon/python/dragon/vm/torch/tensor_pool.py
...@@ -14,6 +14,16 @@ ...@@ -14,6 +14,16 @@
Tensors with the same scope in the pool will be reused by turns, Tensors with the same scope in the pool will be reused by turns,
which speeds up the whole system by reducing the unnecessary deconstructing. 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 from __future__ import absolute_import
...@@ -43,7 +53,7 @@ class TensorPool(object): ...@@ -43,7 +53,7 @@ class TensorPool(object):
self._scope2handle[scope] - 1) self._scope2handle[scope] - 1)
return self._scope2keys[scope].popleft() return self._scope2keys[scope].popleft()
def get(self, scope='runtime'): def get(self, scope='detach'):
handle = self.get_handle(scope) handle = self.get_handle(scope)
return '[TPool]{}/tensor:{}'.format(scope, handle) return '[TPool]{}/tensor:{}'.format(scope, handle)
...@@ -55,7 +65,7 @@ class TensorPool(object): ...@@ -55,7 +65,7 @@ class TensorPool(object):
scope, handle = name[7:].split('/tensor:') scope, handle = name[7:].split('/tensor:')
if not handle.isdigit(): if not handle.isdigit():
raise ValueError('Got illegal name: ' + name + '.\n' 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) self.put_handle(scope, handle)
return True return True
else: return False else: return False
......
Dragon/python/dragon/vm/torch/tensor_uitls.py
...@@ -76,7 +76,7 @@ def from_dragon(tensor, own_storage=False): ...@@ -76,7 +76,7 @@ def from_dragon(tensor, own_storage=False):
""" """
info = dg_tensor_utils.GetTensorInfo(tensor) 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') module = importlib.import_module('dragon.vm.torch.tensor')
th_tensor = getattr(module, type_np2torch(info['dtype']))() th_tensor = getattr(module, type_np2torch(info['dtype']))()
th_tensor._ctx = (info['mem_at'], info['device_id']) 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!