ID-33: new operations(Conv2D + LayerNorm)

server/core/Allocator.cpp

@@ -9,7 +9,7 @@ using namespace std;
 
 Allocator* Allocator::instance = NULL;
 
-#define TMP_BUFF_SIZE 2048
+#define TMP_BUFF_SIZE 4096
 
 void Allocator::startVirtualMode() {
     assert(!instance);

server/core/Blob.cpp

@@ -117,7 +117,7 @@ std::ostream& operator<<(std::ostream& os, const Blob& b) {
         for (int k = 0; k < b.shape.dim3(); ++k){
             for (int i = 0; i < b.shape.rows(); ++i) {
                 for (int j = 0; j < b.shape.cols(); ++j)
-                    os << b(i, j) << " ";
+                    os << b(l, k, i, j) << " ";
                 os << std::endl;
             }
             os << std::endl;

server/core/Layer.cpp

@@ -52,3 +52,65 @@ MSELoss::MSELoss(const std::vector<TensorRef>& args) : mean({0, 1, 2, 3}) {
 MultLayer::MultLayer(const std::vector<TensorRef>& args) {
     result = Tensor(mult, {args[0].get(), args[1].get()});
 }
+
+Conv2DLayer::Conv2DLayer(const Conv2DLayerParameters& params,
+                         const std::vector<TensorRef>& args,
+                         RandomObject* randomInit)
+    : kernel(Blob::constRandomBlob(
+        Shape {{params.outChannels, params.inChannels, params.kernelSize, params.kernelSize}}, 
+        randomInit)
+    ) {
+
+    layerOperationParams.push_back(kernel);
+
+    result = Tensor(conv, {args[0], kernel});
+}
+
+
+
+VarLayer::VarLayer(const AxisParameters& params,
+                   const std::vector<TensorRef>& args)
+    : mean(params.axis), meanMinusOne(params.axis, true), sum(params.axis) {
+    pipeline.reserve(5);
+    TensorRef tensor = args[0];
+    Tensor meanForVar(mean, {tensor});
+    pipeline.push_back(std::move(meanForVar));
+
+    Tensor fillForVar(fill, {tensor, pipeline[0]});
+    pipeline.push_back(std::move(fillForVar));
+
+    Tensor diff(sub, {tensor, pipeline[1]});
+    pipeline.push_back(std::move(diff));
+
+    Tensor square(sqr, {pipeline[2]});
+    pipeline.push_back(std::move(square));
+
+    result = Tensor(meanMinusOne, {pipeline[3]});     
+}
+
+LayerNorm::LayerNorm(const AxisParameters& params,
+                   const std::vector<TensorRef>& args)
+    : varLayer(params, args), mean(params.axis) {
+    pipeline.reserve(6);
+    TensorRef tensor = args[0];
+
+    Tensor mean_(mean, {tensor});
+    pipeline.push_back(std::move(mean_));
+
+    Tensor fill_(fill, {tensor, pipeline[0]});
+    pipeline.push_back(std::move(fill_));
+
+    Tensor diff(sub, {tensor, pipeline[1]});
+    pipeline.push_back(std::move(diff));
+
+    Tensor eps_(eps, {varLayer.result.value()});
+    pipeline.push_back(std::move(eps_));
+
+    Tensor root_(root, {pipeline[3]});
+    pipeline.push_back(std::move(root_));
+
+    Tensor _fill(fill, {tensor, pipeline[4]});
+    pipeline.push_back(std::move(_fill));
+
+    result = Tensor(div, {pipeline[2], pipeline[5]});
+    }

server/core/Layer.h

@@ -54,3 +54,36 @@ class MultLayer: public Layer {
     Multiply mult;
     MultLayer(const std::vector<TensorRef>& args);
 };
+
+class Conv2DLayer: public Layer {
+public:
+    Tensor kernel;
+    Conv2D conv;
+    Conv2DLayer(const Conv2DLayerParameters& params,
+        const std::vector<TensorRef>& args, RandomObject* randomInit = nullptr);
+};
+
+class VarLayer: public Layer {
+public:
+    Mean mean;
+    Mean meanMinusOne;
+    Fill fill;
+    Substract sub;
+    Square sqr;
+    SumAxis sum;
+    VarLayer(const AxisParameters& params,
+        const std::vector<TensorRef>& args);
+};
+
+class LayerNorm: public Layer {
+public:
+    VarLayer varLayer;
+    Mean mean;
+    Fill fill;
+    Substract sub;
+    EPS eps;
+    Root root;
+    Divide div;
+    LayerNorm(const AxisParameters& params,
+        const std::vector<TensorRef>& args);
+};

server/core/LazyBlob.cpp

@@ -111,7 +111,9 @@ class LazyBlobSelfSum final: public LazyBlobReductOperation {
 
 class LazyBlobMean final: public LazyBlobReductOperation {
 public:
-    LazyBlobMean(const LazyBlob &a, std::vector<short> axis): LazyBlobReductOperation(a, axis) {};
+    bool minusOne;
+    LazyBlobMean(const LazyBlob &a, std::vector<short> axis, bool minusOne)
+    : LazyBlobReductOperation(a, axis), minusOne(minusOne) {};
 
     float operator() (std::size_t k, std::size_t l, std::size_t i, std::size_t j) const override {
         float result = 0;
@@ -132,6 +134,8 @@ class LazyBlobMean final: public LazyBlobReductOperation {
             count++;
         });
 
+        if (minusOne)
+            return result / (count - 1);
         return result / count;
     }
 };
@@ -182,6 +186,13 @@ class LazyBlobMult final: public LazyBlobStretchableOperation {
     LazyBlobMult(const LazyBlob &a, const LazyBlob &b): LazyBlobStretchableOperation(a, b, multiply) {};
 };
 
+class LazyBlobDivide final: public LazyBlobStretchableOperation {
+private:
+    static constexpr BinaryTransform divide = [](float x, float y) { return x / y; };
+public:
+    LazyBlobDivide(const LazyBlob &a, const LazyBlob &b): LazyBlobStretchableOperation(a, b, divide) {};
+};
+
 class LazyBlobDot final: public LazyBlobBinaryOperation {
 public:
     LazyBlobDot(const LazyBlob &a, const LazyBlob &b): LazyBlobBinaryOperation(a, b) {};
@@ -222,17 +233,37 @@ class LazyBlobCombine final: public LazyBlobBinaryOperation {
 
 class LazyBlobTranspose final: public LazyBlobUnaryOperation {
 public:
-    LazyBlobTranspose(const LazyBlob &a): LazyBlobUnaryOperation(a) {};
+    bool norm;
+    LazyBlobTranspose(const LazyBlob &a, bool norm = true): LazyBlobUnaryOperation(a), norm(norm) {};
+
+    void initShape() const final override { 
+        if (norm)
+            shape_ = Shape {{a.shape().dim4(), a.shape().dim3(), a.shape().cols(), a.shape().rows()}, a.shape().dimsCount};
+        else
+            shape_ = Shape {{a.shape().dim3(), a.shape().dim4(), a.shape().rows(), a.shape().cols()}, a.shape().dimsCount};
+    }
+
+    float operator() (std::size_t k, std::size_t l, std::size_t i, std::size_t j) const override {
+        if (norm)
+            return a(k, l, j, i);
+        return a(l, k, i, j);
+    }
+};
+
+class LazyBlobReverse final: public LazyBlobUnaryOperation {
+public:
+    LazyBlobReverse(const LazyBlob &a): LazyBlobUnaryOperation(a) {};
 
     void initShape() const final override { 
-        shape_ = Shape {{a.shape().dim4(), a.shape().dim3(), a.shape().cols(), a.shape().rows()}, a.shape().dimsCount};
+        shape_ = Shape {{a.shape().dim4(), a.shape().dim3(), a.shape().rows(), a.shape().cols()}, a.shape().dimsCount};
     }
 
     float operator() (std::size_t k, std::size_t l, std::size_t i, std::size_t j) const override {
-        return a(k, l, j, i);
+        return a(k, l, a.shape().rows() - i - 1, a.shape().cols() - j - 1);
     }
 };
 
+
 class LazyBlobApply: public LazyBlobUnaryOperation {
 private:
     const UnaryTransform operation;
@@ -317,6 +348,11 @@ const LazyBlob& operator * (const LazyBlob &a, const LazyBlob &b) {
     return alloc2<LazyBlobMult>(a, b);
 }
 
+const LazyBlob& operator / (const LazyBlob &a, const LazyBlob &b) {
+    assertStretchable(a, b);
+    return alloc2<LazyBlobDivide>(a, b);
+}
+
 const LazyBlob& LazyBlob::dot(const LazyBlob& a) const {
     assert(shape().cols() == a.shape().rows());
     return alloc2<LazyBlobDot>(*this, a);
@@ -326,14 +362,23 @@ const LazyBlob& LazyBlob::transposed() const {
     return alloc1<LazyBlobTranspose>(*this);
 }
 
+const LazyBlob& LazyBlob::reverseLast2Dims() const {
+    return alloc1<LazyBlobReverse>(*this);
+}
+
+const LazyBlob& LazyBlob::transposeFirst2Dims() const {
+    void* location = Allocator::allocateBytes(sizeof(LazyBlobTranspose));
+    return *(new(location) LazyBlobTranspose(*this, false));
+}
+
 const LazyBlob& LazyBlob::sum(std::vector<short> axis) const {
     void* location = Allocator::allocateBytes(sizeof(LazyBlobSelfSum));
     return *(new(location) LazyBlobSelfSum(*this, axis));
 }
 
-const LazyBlob& LazyBlob::mean(std::vector<short> axis) const {
+const LazyBlob& LazyBlob::mean(std::vector<short> axis, bool minusOne) const {
     void* location = Allocator::allocateBytes(sizeof(LazyBlobMean));
-    return *(new(location) LazyBlobMean(*this, axis));
+    return *(new(location) LazyBlobMean(*this, axis, minusOne));
 }
 
 const LazyBlob& LazyBlob::fill(Shape shape) const {
@@ -415,12 +460,106 @@ Blob& operator *= (Blob& a, const LazyBlob& b) {
     return a;
 }
 
+class LazyBlobConv: public LazyBlob {
+public:
+    const LazyBlob &a, &b;
+    const int size_r, size_c;
+    LazyBlobConv(const LazyBlob &a, const LazyBlob &b): 
+        a(a), b(b), 
+        size_r(b.shape().rows()),
+        size_c(b.shape().cols()) 
+        {};
+
+    float a_get(std::size_t k, std::size_t l, long i, long j) const {
+        if (i < 0 || j < 0 || i >= a.shape().rows() || j >= a.shape().cols())
+            return 0;
+        return a(k, l, i, j);
+    }
+
+    void initShape() const final override { 
+        // TODO: assert
+        shape_ = Shape {{a.shape().dim4(), b.shape().dim4(), a.shape().rows(), a.shape().cols()}, a.shape().dimsCount};
+    }
+
+    float operator() (std::size_t k, std::size_t l, std::size_t i, std::size_t j) const override {
+        float res = 0;
+
+        for (size_t c = 0; c < b.shape().dim3(); ++c) {
+            for (long i1 = 0; i1 < size_r; ++i1) {
+                for (long j1 = 0; j1 < size_c; ++j1) {
+                    res += a_get(k, c, i + i1 - size_r / 2, j + j1 - size_c / 2) * b(l, c, i1, j1);
+                }
+            }
+        }
+        return res;
+    }
+};
+
+LazyBlob&  conv(const LazyBlob &a, const LazyBlob &b) {
+    assert(a.shape().dim3() == b.shape().dim3());
+    void* location = Allocator::allocateBytes(sizeof(LazyBlobConv));
+    return *(new(location) LazyBlobConv(a, b));
+}
+
+class LazyBlobConvI: public LazyBlob {
+public:
+    const LazyBlob &a, &b;
+    const size_t kernelSize, index;
+    LazyBlobConvI(const LazyBlob &a, const LazyBlob &b, size_t kernelSize, size_t i): 
+        a(a), b(b), kernelSize(kernelSize), index(i) {};
+
+    float a_get(std::size_t k, std::size_t l, long i, long j) const {
+        if (i < 0 || j < 0 || i >= a.shape().rows() || j >= a.shape().cols())
+            return 0;
+        return a(k, l, i, j);
+    }
+
+    void initShape() const final override { 
+        shape_ = Shape {{b.shape().dim3(), a.shape().dim3(), kernelSize, kernelSize}, a.shape().dimsCount};
+    }
+
+    float operator() (std::size_t k, std::size_t l, std::size_t i, std::size_t j) const override {
+        float res = 0;
+        for (long i1 = 0; i1 < a.shape().rows(); ++i1) {
+            for (long j1 = 0; j1 < a.shape().cols(); ++j1) {
+                res += a_get(index, l, i1 + i - kernelSize / 2, j1 + j - kernelSize / 2) * b(index, k, i1, j1);
+            }
+        }
+        return res;
+    }
+};
+
+LazyBlob&  conv_i(const LazyBlob &a, const LazyBlob &b, size_t kernelSize, size_t i) {
+    assert(a.shape().dim4() == b.shape().dim4());
+    void* location = Allocator::allocateBytes(sizeof(LazyBlobConvI));
+    return *(new(location) LazyBlobConvI(a, b, kernelSize, i));
+}
+
+class LazyBlobZero: public LazyBlob {
+public:
+    const Shape myShape;
+    LazyBlobZero(const Shape& shape): myShape(shape) {};
+
+    void initShape() const final override {
+        shape_ = myShape;
+    }
+
+    float operator() (std::size_t k, std::size_t l, std::size_t i, std::size_t j) const override {
+        return 0;
+    }
+};
+
+LazyBlob&  zeroBlob(const Shape& shape) {
+    void* location = Allocator::allocateBytes(sizeof(LazyBlobZero));
+    return *(new(location) LazyBlobZero(shape));
+}
+
 std::ostream& operator<<(std::ostream& os, const LazyBlob &b) {
     for (int l = 0; l < b.shape().dim4(); ++l) {
         for (int k = 0; k < b.shape().dim3(); ++k){
             for (int i = 0; i < b.shape().rows(); ++i) {
                 for (int j = 0; j < b.shape().cols(); ++j)
-                    os << b(i, j) << " ";
+                    os << b(l, k, i, j) << " ";
                 os << std::endl;
             }
             os << std::endl;

server/core/LazyBlob.h

@@ -28,7 +28,9 @@ class LazyBlob {
     const LazyBlob& transposed() const;
     const LazyBlob& applying(const UnaryTransform t) const;
     const LazyBlob& sum(std::vector<short> axis) const;
-    const LazyBlob& mean(std::vector<short> axis) const;
+    const LazyBlob& mean(std::vector<short> axis, bool minusOne = false) const;
+    const LazyBlob& reverseLast2Dims() const;
+    const LazyBlob& transposeFirst2Dims() const;
 
     /// To repeat some dimensions several times
     /// - Parameter shape: the size we want to get
@@ -39,6 +41,7 @@ class LazyBlob {
     friend const LazyBlob& operator - (const LazyBlob &a);
     /// ELEMENT-WISE
     friend const LazyBlob& operator * (const LazyBlob &a, const LazyBlob &b);
+    friend const LazyBlob& operator / (const LazyBlob &a, const LazyBlob &b);
     /// MATRIX
     friend const LazyBlob& operator & (const LazyBlob &a, const LazyBlob &b);
 
@@ -50,6 +53,10 @@ class LazyBlob {
     friend std::ostream& operator<<(std::ostream& os, const LazyBlob& b);
 };
 
+LazyBlob&  conv(const LazyBlob &a, const LazyBlob &b);
+LazyBlob&  conv_i(const LazyBlob &a, const LazyBlob &b, std::size_t kernelSize, std::size_t i);
+LazyBlob& zeroBlob(const Shape& shape);
+
 class LazyBlobView final: public LazyBlob {
 private:
     const Blob &ref;