Update layers.cpp

examples/layers.cpp

@@ -1,28 +1,34 @@
 #include <CaberNet/CaberNet.h>
 
-struct Encoder : public net::base::Model {
-    public:
-    Encoder() {
-        layers = {
-            new net::layer::Linear(64, 32),
-            new net::layer::ReLU(),
-            new net::layer::Linear(32, 16),
-            new net::layer::ReLU(),
-        };
-    }
+struct Autoencoder : public net::Model<Autoencoder> {
+
+    Autoencoder() = default;
+
+    net::layer::Sequence encoder {
+        net::layer::Linear(784, 128, net::initializer::He),
+        net::layer::ReLU(),
+        net::layer::Linear(128, 64, net::initializer::He),
+    };
+
+    net::layer::Sequence decoder {
+        net::layer::Linear(64, 128, net::initializer::He),
+        net::layer::ReLU(),
+        net::layer::Linear(128, 784, net::initializer::He),
+        net::layer::LogSoftmax(1/*axis*/)
+    };
 
     net::Tensor forward(net::Tensor x) {
-        return layers.forward(x);
+        x = encoder(x);
+        x = decoder(x);
+        return x;
     }
 
-    net::layer::Sequence layers;
 };
 
-
 int main() {
-    Encoder encoder;
-    net::Tensor input({256, 64}); input.fill(net::initializer::He);
-    net::Tensor output = encoder.forward(input);
-    output.perform();
-    for(auto element : output) std::cout << element;
-}
+    Autoencoder model;
+    net::Tensor x({1, 784}); x.fill(net::initializer::He);
+    net::Tensor y = model(x);
+    y.perform();
+    std::cout << y;
+}