diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index adea88ef..9554a42c 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -51,6 +51,7 @@ jobs:
           sudo apt install -y g++-11
           sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 90
           sudo update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-11 90
+          sudo apt install libzip-dev
 
       - name: install-cpprest
         run: sudo apt install libcpprest-dev
diff --git a/.gitignore b/.gitignore
index 48b197a5..f966ca0c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,6 +1,10 @@
 Build/
 build/
 
+# Large-files
+zip-examples
+model_data
+
 # IDE
 .vscode
 .idea
@@ -10,6 +14,7 @@ build/
 .DS_Store
 .AppleDouble
 .LSOverride
+*.entitlements
 
 # Javascript
 package-lock.json
diff --git a/client/scripts/requests.js b/client/scripts/requests.js
index 35124124..c8049021 100644
--- a/client/scripts/requests.js
+++ b/client/scripts/requests.js
@@ -124,29 +124,60 @@ async function deleteConnection(sending_object) {
     return response
 }
 
-function trainRequest() {
-    if (!train_data) {
-        errorNotification("No training data was set.")
-    } else {
-        fetch(`http://${py_server_address}/train/${user_id}/${model_id}/0`, {
-            method: "PUT",
-            mode: "cors",
-            headers: {"Content-Type": "text/csv"},
-            body: train_data,
-        }).then(response => {
-            showBuildNotification(response.ok)
-            onTrainShowPredict(response.ok)
-            if (response.ok) {
-                setModelView("success")
-            } else {
-                setModelView("error")
-            }
+function uploadRequest() {
+    if (data_upload.files.length == 0) return
+    const file = data_upload.files[0]
+
+    fetch(`http://${py_server_address}/${user_id}/${model_id}`, {
+        method: "PATCH",
+        mode: "cors",
+        body: file,
+    }).then(response => {
+        if (!response.ok) {
+            Swal.fire({
+                position: "top-end",
+                icon: "error",
+                title: "Failed to upload data",
+                showConfirmButton: false,
+                timer: 1500,
+            })
+            console.error(`Failed to upload data for ${file.name}`)
+            return
+        }
+        Swal.fire({
+            position: "top-end",
+            icon: "success",
+            title: "Successfully uploaded",
+            showConfirmButton: false,
+            timer: 1500,
         })
-    }
+    })
+    setModelView("irrelevant")
+    // allow user to press a train button from now on
+    button_wrapper = document.getElementById("train-button")
+    button_wrapper.getElementsByTy
+    train_button = button_wrapper.children[0]
+    button_wrapper.removeAttribute("disabled")
+    train_button.removeAttribute("disabled")
+}
+
+function trainRequest() {
+    fetch(`http://${py_server_address}/train/${user_id}/${model_id}/0`, {
+        method: "PUT",
+        mode: "cors",
+    }).then(response => {
+        showBuildNotification(response.ok)
+        onTrainShowPredict(response.ok)
+        if (response.ok) {
+            setModelView("success")
+        } else {
+            setModelView("error")
+        }
+    })
 }
 
 async function predictRequest() {
-    if (csv_predict.files.length == 0) {
+    if (predict_button.files.length == 0) {
         errorNotification("Empty predict file.")
         return
     }
@@ -159,25 +190,34 @@ async function predictRequest() {
             showConfirmButton: true,
         })
     }
-    const file = csv_predict.files[0]
-    const text = await file.text()
-    const response = await fetch(
+    const file = predict_button.files[0]
+    let response = await fetch(
         `http://${py_server_address}/predict/${user_id}/${model_id}`,
         {
             method: "PUT",
             mode: "cors",
-            headers: {"Content-Type": "text/csv"},
-            body: text,
+            body: file,
         },
     )
+    if (!response.ok) {
+        Swal.fire("Error!", "Failed to upload the png image", "error")
+        console.error(
+            `Failed to upload the png with ${response.statusText}: ${responseJson.error}`,
+        )
+        return
+    }
+    response = await fetch(
+        `http://${py_server_address}/predict/${user_id}/${model_id}`,
+        {method: "GET", mode: "cors"},
+    )
     const responseJson = await response.json()
-    hideResult() // hide previous predict result
-    if (response.ok) onPredictShowResult(responseJson)
-    else {
-        Swal.fire("Error!", "Server is not responding now.", "error")
-        errorNotification("Failed to predict.\n" + responseJson.error)
-        console.log(
-            `Predict failed with ${response.statusText}: ${responseJson.error}`,
+    if (!response.ok) {
+        Swal.fire("Error!", "Failed to predict", "error")
+        console.error(
+            `Failed to predict with with ${response.statusText}: ${responseJson.error}`,
         )
+        return
     }
+    hideResult() // hide previous predict result
+    onPredictShowResult(responseJson)
 }
diff --git a/client/templates/main.html b/client/templates/main.html
index a8122627..78bffe8c 100644
--- a/client/templates/main.html
+++ b/client/templates/main.html
@@ -113,7 +113,7 @@ <h2 onmousedown="return false">MLCraft</h2>
 
         <label class="button-wrapper">
           Upload
-          <input type="file" id="csv-upload-button" accept=".csv" />
+          <input type="file" id="dataset-upload-button" accept=".csv,.zip" />
         </label>
 
         <label class="button-wrapper" id="train-button" disabled>
@@ -123,7 +123,7 @@ <h2 onmousedown="return false">MLCraft</h2>
 
         <label class="button-wrapper" disabled>
           Predict
-          <input type="file" accept=".csv" id="csv-predict-button" disabled />
+          <input type="file" accept=".png" id="predict-button" disabled />
         </label>
         <div class="predict-result-wrapper" id="predict-result-wrapper">
           <output class="form-control-res" id="res-value-out" name="res-value" for="x-value y-value"></output>
@@ -329,6 +329,7 @@ <h3></h3>
           deleteLayer(layer_to_delete).then(response => {
             if (response.ok) {
               super.removeNodeId(id)
+              updateConnectionsInLayers(id.split("-")[1])
             } else {
               errorNotification("Failed to delete layer.\n Server is not responding now.")
             }
@@ -395,15 +396,15 @@ <h3></h3>
             target: "Target",
           }
           const parametersMap = {
-            linear: "bias=0;inFeatures=1;outFeatures=2",
+            linear: "bias=0;inFeatures=2;outFeatures=1",
             relu: "",
             pooling: "",
             layernorm: "axis=1",
             conv2d: "KernelSize=1;inChannels=1;outChannels=1",
-            data: "width=2",
-            output: "width=1",
-            loss: "width=1",
-            target: "width=1",
+            data: "shape=[2]",
+            output: "",
+            loss: "",
+            target: "shape=[1]",
           }
           const layer = {
             type: nameMap[name],
@@ -569,7 +570,7 @@ <h3></h3>
           if (response.ok) {
             editor.updateNodeDataFromId(connection.input_id, {
               DBID: data2.DBID,
-              Connections: [].concat(`${data1.DBID};${connection.output_id}`, data2.Connections), // "db_id;drawflow_id"
+              Connections: data2.Connections.concat(`${data1.DBID};${connection.output_id}`), // "db_id;drawflow_id"
               Parameters: data2.Parameters,
             })
           } else {
@@ -1011,6 +1012,8 @@ <h3></h3>
         metricsPage.appendChild(imageElement)
         URL.revokeObjectURL(imageUrl)
       }
+
+      setInterval(getMetricsFromServer, 10000)
     </script>
     <!--Метрики-->
 
@@ -1113,12 +1116,12 @@ <h3></h3>
 
     <!--JSON requests handling-->
     <script>
-      const csv_predict = document.getElementById("csv-predict-button")
-      csv_predict.onclick = () => {
-        csv_predict.value = null
+      const predict_button = document.getElementById("predict-button")
+      predict_button.onclick = () => {
+        predict_button.value = null
       }
 
-      csv_predict.addEventListener("input", predictRequest)
+      predict_button.addEventListener("input", predictRequest)
 
       function buildJsonFormData(form) {
         const jsonFormData = {}
@@ -1282,26 +1285,14 @@ <h3></h3>
       const showBuildNotification = build_status => {
         build_status ? buildSuccessful() : buildUnsuccessful()
       }
-      async function csvUploaded() {
-        if (csv_upload.files.length == 0) return
-        const file = csv_upload.files[0]
-        const text = await file.text()
-        // do whatever with the `text`
-        setModelView("irrelevant")
-        console.log(`Got goods: ${text}`)
-        train_data = text
-        // allow user to press a train button from now on
-        button_wrapper = document.getElementById("train-button")
-        button_wrapper.getElementsByTy
-        train_button = button_wrapper.children[0]
-        button_wrapper.removeAttribute("disabled")
-        train_button.removeAttribute("disabled")
+      const data_upload = document.getElementById("dataset-upload-button")
+      data_upload.onclick = () => {
+        data_upload.value = null
       }
-      const csv_upload = document.getElementById("csv-upload-button")
-      csv_upload.addEventListener("change", csvUploaded)
+      data_upload.addEventListener("change", uploadRequest)
 
       function onTrainShowPredict(training_successful) {
-        const predict_button = document.getElementById("csv-predict-button")
+        const predict_button = document.getElementById("predict-button")
         hideResult() // hide result when new training results arrive
         if (training_successful) {
           predict_button.removeAttribute("disabled")
@@ -1412,6 +1403,7 @@ <h3></h3>
 
         // Input selection
         writeIdOnConnections(input_selection_layer_id)
+
         let sortable_list = getSortableInputList(node)
         let child = sortable_list.lastElementChild
 
@@ -1440,8 +1432,9 @@ <h3></h3>
             new_input_order = []
             for (let i in sortable_list.childNodes) {
               try {
-                if (sortable_list.childNodes[i].textContent.includes("input")) {
-                  new_input_order.push(Number(sortable_list.childNodes[i].textContent.split(" ")[3]))
+                if (sortable_list.childNodes[i].textContent.length > 0) {
+                  console.log(sortable_list.childNodes[i].textContent)
+                  new_input_order.push(Number(sortable_list.childNodes[i].textContent))
                 }
               } catch (e) {
                 continue
@@ -1450,7 +1443,7 @@ <h3></h3>
             matchConnectionsAfterSelection(new_input_order)
           },
         })
-        document.getElementById("layer-id").innerHTML = "ID: " + layer_dbid.toString()
+
         document.getElementById("layer-data").style.display = "block"
         setupLayerParameters(node, correct_id)
       }
@@ -1651,6 +1644,25 @@ <h3></h3>
           }
         }
       }
+
+      function updateConnectionsInLayers(deleted_id) {
+        let nodes = document.getElementsByClassName("drawflow-node")
+        for (let i of nodes) {
+          const data1 = editor.getNodeFromId(i.id.split("-")[1]).data
+          let updatedConnections = []
+          for (let connection of data1.Connections) {
+            if (connection.split(";")[1] == deleted_id) {
+              continue
+            }
+            updatedConnections.unshift(connection)
+          }
+          editor.updateNodeDataFromId(i.id.split("-")[1], {
+            DBID: data1.DBID,
+            Connections: updatedConnections,
+            Parameters: data1.Parameters,
+          })
+        }
+      }
     </script>
 
     <!--Терминал-->
diff --git a/documentation/png-examples/Ours.png b/documentation/png-examples/Ours.png
new file mode 100644
index 00000000..3e4a553c
Binary files /dev/null and b/documentation/png-examples/Ours.png differ
diff --git a/documentation/png-examples/Vectors.png b/documentation/png-examples/Vectors.png
new file mode 100644
index 00000000..6366e886
Binary files /dev/null and b/documentation/png-examples/Vectors.png differ
diff --git a/py_server/mlcraft/__init__.py b/py_server/mlcraft/__init__.py
index 912740c6..2df1fdb3 100644
--- a/py_server/mlcraft/__init__.py
+++ b/py_server/mlcraft/__init__.py
@@ -63,6 +63,7 @@ def make_app(config=None):
     app.register_error_handler(KeyError, key_error)  # no json field
     app.register_error_handler(ValueError, value_error)  # json wrong type
     app.register_error_handler(HTTPException, http_error)
+    app.register_error_handler(TimeoutError, timeout_error)
     app.register_error_handler(ConnectTimeout, timeout_error)
     app.register_error_handler(ConnectionError, connection_error)
 
diff --git a/py_server/mlcraft/check_dimensions.py b/py_server/mlcraft/check_dimensions.py
index 772e928d..5147ff80 100644
--- a/py_server/mlcraft/check_dimensions.py
+++ b/py_server/mlcraft/check_dimensions.py
@@ -6,8 +6,8 @@
 
 
 class Data2dChecker:
-    def __init__(self, width):
-        self.output_shape = [-1, width]
+    def __init__(self, shape: list[int]):
+        self.output_shape = [-1] + shape
 
     def __call__(self):
         return True
@@ -60,7 +60,7 @@ def __call__(self, input_shape: list[int], targets_shape: list[int]):
 def create_checker(layer: dict):
     match layer["type"]:
         case "Data" | "Target":
-            return Data2dChecker(layer["parameters"]["width"])
+            return Data2dChecker(layer["parameters"]["shape"])
         case "Linear":
             return LinearChecker(
                 layer["parameters"]["inFeatures"], layer["parameters"]["outFeatures"]
diff --git a/py_server/mlcraft/dataset.py b/py_server/mlcraft/dataset.py
deleted file mode 100644
index 79fcc2ab..00000000
--- a/py_server/mlcraft/dataset.py
+++ /dev/null
@@ -1,46 +0,0 @@
-import csv
-from http import HTTPStatus
-from io import StringIO
-from .errors import Error
-
-
-def get_layer(type: str, model: dict) -> dict:
-    layer = next(filter(lambda l: l["type"] == type, model["layers"]), None)
-    if layer is None:
-        raise Error(f"No {type} block in model", HTTPStatus.PRECONDITION_FAILED)
-    return layer
-
-
-def extract_train_data(bytes: bytes, model: dict) -> dict:
-    data_layer = get_layer("Data", model)
-    target_layer = get_layer("Target", model)
-
-    reader = csv.reader(StringIO(bytes.decode()))
-
-    data: list[float] = []
-    target: list[float] = []
-    for row in reader:
-        if (
-            int(data_layer["parameters"]["width"]) != len(row) - 1
-        ):  # columns = features + 1 for target
-            raise Error("Input csv column count doesn't match data's feature count")
-
-        data.extend(map(float, row[:-1]))
-        target.append(float(row[-1]))
-
-    return {data_layer["id"]: data, target_layer["id"]: target}
-
-
-def extract_predict_data(bytes: bytes, model: dict) -> dict:
-    data_layer = get_layer("Data", model)
-    reader = csv.reader(StringIO(bytes.decode()))
-    row = next(reader, None)
-
-    if row is None:
-        raise Error("Bad csv format")
-    if int(data_layer["parameters"]["width"]) != len(
-        row
-    ):  # predict request - one row with data only
-        raise Error("Input csv column count doesn't match data's feature count")
-
-    return {data_layer["id"]: list(map(float, row))}
diff --git a/py_server/mlcraft/errors.py b/py_server/mlcraft/errors.py
index f9e3575a..86b912ec 100644
--- a/py_server/mlcraft/errors.py
+++ b/py_server/mlcraft/errors.py
@@ -44,7 +44,7 @@ def http_error(e: HTTPException):
     return {"error": str(e)}, HTTPStatus.BAD_REQUEST
 
 
-def timeout_error(e: ConnectTimeout):
+def timeout_error(e: TimeoutError | ConnectTimeout):
     return {"error": "Request to c++ server timeout"}, HTTPStatus.INTERNAL_SERVER_ERROR
 
 
diff --git a/py_server/mlcraft/server.py b/py_server/mlcraft/server.py
index d81137d3..ed9dc61a 100644
--- a/py_server/mlcraft/server.py
+++ b/py_server/mlcraft/server.py
@@ -1,6 +1,3 @@
-from json import dumps
-from sqlite3 import IntegrityError
-import datetime
 import requests
 from http import HTTPStatus
 from flask import Blueprint, request, current_app, send_file
@@ -12,19 +9,21 @@
     is_valid_model,
     convert_model,
     plot_metrics,
-    delete_file,
 )
 from .check_dimensions import assert_dimensions_match
 
 from .errors import Error
 
 from .db import sql_worker
-from .dataset import extract_predict_data, extract_train_data
 
 
 app = Blueprint("make a better name", __name__)
 
 
+def cpp_url(method: str):
+    return current_app.config["CPP_SERVER"] + "/" + method
+
+
 @app.route("/user", methods=["POST"])
 def add_user():
     json_data = request.json
@@ -63,6 +62,14 @@ def model(user_id: int, model_id: int):
             d: dict[str, str | None] = defaultdict(lambda: None, **request.json)  # type: ignore
             sql_worker.update_model(model_id, d["name"], d["raw"])
             return "", HTTPStatus.OK
+        case "PATCH":
+            response = requests.post(
+                cpp_url(f"upload_data/{model_id}/0"),
+                data=request.data,
+                headers={"Content-Type": request.content_type},
+                timeout=10,
+            )
+            return "", response.status_code
         case "DELETE":
             sql_worker.delete_model(model_id)
             return "", HTTPStatus.OK
@@ -164,13 +171,11 @@ def train_model(
         raise Error("Invalid model found", HTTPStatus.NOT_ACCEPTABLE)
     assert_dimensions_match(model["layers"])
 
-    dataset = extract_train_data(request.data, model)
-
     convert_model(model)
-    model = {"graph": model, "dataset": dataset}
+    model = {"graph": model}
 
     response = requests.post(
-        current_app.config["CPP_SERVER"] + f"/train/{user_id}/{model_id}",
+        cpp_url(f"train/{user_id}/{model_id}"),
         json=model,
         timeout=3,
     )
@@ -179,25 +184,25 @@ def train_model(
     return response.text, response.status_code
 
 
-@app.route("/predict/<int:user_id>/<int:model_id>", methods=["PUT"])
+@app.route("/predict/<int:user_id>/<int:model_id>", methods=["GET", "PUT"])
 def predict(user_id: int, model_id: int):
+    """PUT method is for uploading the png, GET method is for receiving the result"""
     sql_worker.verify_access(user_id, model_id)
 
-    if not request.data:
-        raise Error("No csv data provided")
-
-    model = sql_worker.get_graph_elements(model_id)
-    convert_model_parameters(model)
-    json_data = extract_predict_data(request.data, model)
-
     if not sql_worker.is_model_trained(model_id):
         raise Error("Not trained", HTTPStatus.PRECONDITION_FAILED)
 
-    response = requests.post(
-        current_app.config["CPP_SERVER"] + f"/predict/{user_id}/{model_id}",
-        json=json_data,
-        timeout=3,
-    )
+    match request.method:
+        case "GET":
+            response = requests.put(
+                cpp_url(f"predict/{model_id}"),
+            )
+        case "PUT":
+            response = requests.post(
+                cpp_url(f"upload_data/{model_id}/1"),
+                data=request.data,
+                headers={"Content-Type": "image/png"},
+            )
 
     return response.text, response.status_code
 
@@ -265,5 +270,5 @@ def get_plots(user_id: int, model_id: int):
     current_dir = os.getcwd()
     print(current_dir)
     response = send_file(os.path.join(current_dir, "images", plot_path))
-    delete_file(os.path.join(current_dir, "images", plot_path))
+    # delete_file(os.path.join(current_dir, "images", plot_path))
     return response
diff --git a/py_server/mlcraft/static/swagger.yaml b/py_server/mlcraft/static/swagger.yaml
index 6b891f9c..ed48beda 100644
--- a/py_server/mlcraft/static/swagger.yaml
+++ b/py_server/mlcraft/static/swagger.yaml
@@ -479,18 +479,11 @@ paths:
         - $ref: '#/components/parameters/ModelId'
         - name: safe
           in: path
-          description: flag, if 1, then the training process will be started even if the model is already trained
+          description: flag, if 0, then the training process will be started even if the model is already trained
           required: true
           schema:
             type: integer
             format: int64
-      requestBody:
-        description: "Data for training"
-        required: true
-        content:
-          text/csv:
-            schema:
-              $ref: '#/components/schemas/TrainData'
       responses:
         '200':
           description: "Model is trained"
@@ -532,13 +525,6 @@ paths:
       parameters:
         - $ref: '#/components/parameters/UserId'
         - $ref: '#/components/parameters/ModelId'
-      requestBody:
-        description: "Data for predictions"
-        required: true
-        content:
-          text/csv:
-            schema:
-              $ref: '#/components/schemas/PredictData'
       responses:
         '200':
           description: "Prediction was successful"
@@ -886,21 +872,6 @@ components:
           items:
             type: number
             example: [0.0, 0.0, 0.0, 1.0]
-    TrainData:
-      type: string
-      description: |-
-        Last column is for the target value 
-        and all other columns are features values
-      example: |-
-        0,0,0
-        0,1,1
-        1,0,1
-        1,1,1
-    PredictData:
-      type: string
-      description: Only ONE row for now, a column for every feature
-      example: |-
-        1,1
     PredictResult:
       type: array
       items:
diff --git a/py_server/tests/test_check_dimensions.py b/py_server/tests/test_check_dimensions.py
index 705f5e36..35cdbf20 100644
--- a/py_server/tests/test_check_dimensions.py
+++ b/py_server/tests/test_check_dimensions.py
@@ -5,8 +5,8 @@
 
 def test_correct_simple():
     layers = [
-        {"id": 0, "type": "Target", "parameters": {"width": 1}, "parents": []},
-        {"id": 1, "type": "Data", "parameters": {"width": 2}, "parents": []},
+        {"id": 0, "type": "Target", "parameters": {"shape": [1]}, "parents": []},
+        {"id": 1, "type": "Data", "parameters": {"shape": [2]}, "parents": []},
         {
             "id": 2,
             "type": "Linear",
@@ -23,8 +23,8 @@ def test_correct_simple():
 
 def test_correct_harder():
     layers = [
-        {"id": 0, "type": "Target", "parameters": {"width": 1}, "parents": []},
-        {"id": 1, "type": "Data", "parameters": {"width": 2}, "parents": []},
+        {"id": 0, "type": "Target", "parameters": {"shape": [1]}, "parents": []},
+        {"id": 1, "type": "Data", "parameters": {"shape": [2]}, "parents": []},
         {
             "id": 2,
             "type": "Linear",
@@ -55,8 +55,8 @@ def test_correct_harder():
 
 def test_mismatch_simple():
     layers = [
-        {"id": 0, "type": "Target", "parameters": {"width": 1}, "parents": []},
-        {"id": 1, "type": "Data", "parameters": {"width": 2}, "parents": []},
+        {"id": 0, "type": "Target", "parameters": {"shape": [1]}, "parents": []},
+        {"id": 1, "type": "Data", "parameters": {"shape": [2]}, "parents": []},
         {
             "id": 2,
             "type": "Linear",
@@ -77,8 +77,8 @@ def test_mismatch_simple():
 
 def test_mismatch_harder():
     layers = [
-        {"id": 0, "type": "Target", "parameters": {"width": 1}, "parents": []},
-        {"id": 1, "type": "Data", "parameters": {"width": 2}, "parents": []},
+        {"id": 0, "type": "Target", "parameters": {"shape": [1]}, "parents": []},
+        {"id": 1, "type": "Data", "parameters": {"shape": [2]}, "parents": []},
         {
             "id": 2,
             "type": "Linear",
diff --git a/server/Makefile b/server/Makefile
index 02af06b4..b3f865fa 100644
--- a/server/Makefile
+++ b/server/Makefile
@@ -126,7 +126,7 @@ serve: $(server) | $(dirTrain) $(dirPredict)
 	$< $(host) $(port)
 
 $(server): $(serverMain) $(serverObjects) $(coreArchive) | checkboost $(dirBin)
-	$(LINK.cpp) -L $(boost)/lib $^ -o $@ -lX11 -lpng -lcpprest -lssl -lcrypto
+	$(LINK.cpp) -L $(boost)/lib $^ -o $@ -lX11 -lpng -lcpprest -lssl -lcrypto -lzip
 
 
 ### TESTS ###
diff --git a/server/api/DataLoader.cpp b/server/api/DataLoader.cpp
index 9c643ed4..765fc7f0 100644
--- a/server/api/DataLoader.cpp
+++ b/server/api/DataLoader.cpp
@@ -38,39 +38,38 @@ std::pair<Blob, std::vector<float>> DataLoader::operator[](std::size_t index) co
 }
 
 std::size_t DataLoader::size() const {
+
     return loader->size();
 }
 
+std::size_t DataLoader::batch_count() const {
+    return (loader->size() + batch_size - 1) / batch_size;
+}
+
 void DataLoader::add_data(const DataLoader& other, int index) {
     loader->add_data(other.loader, index);
 }
 
-std::pair<std::vector<float>, std::vector<float>> DataLoader::get_raw(std::size_t index) const { // batch_size lines from index
-    if (index >= loader->size()) {
+std::pair<std::vector<float>, std::vector<float>> DataLoader::get_raw(std::size_t batch_index) const { // batch_size lines from index
+    if (batch_index >= batch_count()) {
         throw std::out_of_range("Index out of range");
     }
     std::vector<float> data;
     std::vector<float> res(batch_size, 0);
-    Shape shape = loader->get_appropriate_shape(rearrangement[index], batch_size);
-    auto dims = shape.getDims();
-    int data_size = 1;
-    for (int i = 0; i < dims.size(); ++i) {
-        data_size *= dims[i];
-    }
-    data.resize(data_size, 0);
+    Shape shape = loader->get_appropriate_shape(rearrangement[batch_index], batch_size);
+    data.resize(shape.size(), 0);
     int cur_data = 0;
-    for (int i = index; i < index + batch_size; ++i) {
+    for (int i = batch_size * batch_index; i < batch_size * (batch_index + 1); ++i) {
         if (i >= loader->size()) {
             break;
         }
         auto line = loader->get_raw(rearrangement[i]);
-        res[i - index] = line.second;
+        res[i - batch_size * batch_index] = line.second;
         for (int j = 0; j < line.first.size(); ++j) {
-            data[cur_data] = line.first[j];
-            cur_data++;
+            data[cur_data++] = line.first[j];
         }
     }
-    return {data, res};
+    return {std::move(data), std::move(res)};
 }
 
 void DataLoader::shuffle() {
diff --git a/server/api/DataLoader.h b/server/api/DataLoader.h
index a10582cb..b87d1931 100644
--- a/server/api/DataLoader.h
+++ b/server/api/DataLoader.h
@@ -18,6 +18,7 @@ class DataLoader {
     std::pair<Blob, std::vector<float>> operator[](std::size_t index) const;
     void add_data(const DataLoader& other, int index);
     std::size_t size() const;
+    std::size_t batch_count() const;
     std::pair<std::vector<float>, std::vector<float>> get_raw(std::size_t index) const;
     void shuffle();
 };
diff --git a/server/api/GraphBuilder.cpp b/server/api/GraphBuilder.cpp
index afda854d..def35368 100644
--- a/server/api/GraphBuilder.cpp
+++ b/server/api/GraphBuilder.cpp
@@ -2,12 +2,7 @@
 #include "Allocator.h"
 
 void Graph::OverviewLayers(const crow::json::rvalue& layers, 
-                           const std::vector<std::vector<float>>& data,
-                           std::unordered_map<int, crow::json::rvalue>& layer_dicts,
-                           std::unordered_map<int, std::vector<float>>& data_dicts) {
-    std::vector<float> instances;
-    std::vector<float> answers;
-    ParseCsvData(data, instances, answers);
+                           std::unordered_map<int, crow::json::rvalue>& layer_dicts) {
     for (auto& layer : layers) {
         CHECK_HAS_FIELD(layer, "id");
         CHECK_HAS_FIELD(layer, "type");
@@ -16,14 +11,6 @@ void Graph::OverviewLayers(const crow::json::rvalue& layers,
 
         layerTypes_[id] = type;
         layer_dicts[id] = layer;
-
-        if (type == "Data") {
-            data_dicts[id] = instances; // All data goes for all data layers. May be should be splitted by them
-            dataIds_.push_back(id);
-        }
-        if (type == "Target") {
-            data_dicts[id] = answers;
-        }
     }
 }
 
@@ -88,9 +75,9 @@ void Graph::TopologySort(std::unordered_map<int, std::vector<int>>& edges,
 }
 
 void Graph::Initialize(crow::json::rvalue modelJson,
-             const std::vector<std::vector<float>>& data,
-             RandomObject* randomInit,
-             OptimizerBase& SGD) {
+                       RandomObject* randomInit,
+                       OptimizerBase& SGD,
+                       size_t batch_size) {
     Allocator::startVirtualMode();
     CHECK_HAS_FIELD(modelJson, "graph");
     CHECK_HAS_FIELD(modelJson["graph"], "layers");
@@ -101,8 +88,7 @@ void Graph::Initialize(crow::json::rvalue modelJson,
 
     // Parse Jsons into dicts of Jsons
     std::unordered_map<int, crow::json::rvalue> layerDicts;
-    std::unordered_map<int, std::vector<float>> dataDicts;
-    OverviewLayers(layersJson, data, layerDicts, dataDicts);
+    OverviewLayers(layersJson, layerDicts);
 
     std::unordered_map<int, std::vector<int>> straightEdges, reversedEdges;
     std::unordered_set<int> entryNodes;
@@ -133,12 +119,8 @@ void Graph::Initialize(crow::json::rvalue modelJson,
                 targetsIds_.push_back(layer_id);
             }
             CHECK_HAS_FIELD(layerDicts[layer_id], "parameters");
-            auto params = ParseData2d(layerDicts[layer_id]["parameters"]);
-            if (dataDicts[layer_id].size() % params.width != 0) {
-                throw std::invalid_argument("Sizes mismatch!");
-            }
-            params.height = dataDicts[layer_id].size() / params.width;
-            layers_.emplace(layer_id, new Data2dLayer{params, dataDicts[layer_id]});
+            Shape shape = ParseData(layerDicts[layer_id]["parameters"]);
+            layers_.emplace(layer_id, new DataLayer{shape, batch_size});
         } else if (type == "Output") {
             for (auto prevLayerId : reversedEdges[layer_id]) {
                 lastPredictIds_.push_back(prevLayerId);
@@ -152,18 +134,28 @@ void Graph::Initialize(crow::json::rvalue modelJson,
     }
 }
 
-void Graph::ChangeInputData(std::vector<float> data) {
+void Graph::ChangeLayersData(std::vector<float> data, BaseLayerType type) {
     // All data goes to every data layer. Should be changed?
-    for (int id : dataIds_) {
-        Data2dLayer* layer = reinterpret_cast<Data2dLayer*>(layers_[id]);
+    std::vector<int>* layers = nullptr;
+    if (type == BaseLayerType::Data) {
+        layers = &dataIds_;
+    } else if (type == BaseLayerType::Targets) {
+        layers = &targetsIds_;
+    } else {
+        throw std::invalid_argument("Can change data only in 'Data' or 'Target' layers");
+    }
+    for (int id : *layers) {
+        DataLayer* layer = reinterpret_cast<DataLayer*>(layers_[id]);
+
+        Shape expected_shape = layer->result->output->shape;
+        size_t sample_size = expected_shape.stride(4 - expected_shape.dimsCount);
 
-        size_t width = layer->result->output->shape.cols();
-        if (data.size() % width != 0) {
+        if (data.size() % sample_size != 0 || data.size() > expected_shape.size()) {
+            std::cerr << data.size() << " " << sample_size << " " << expected_shape.size() << std::endl;
             throw std::invalid_argument("Sizes mismatch!");
         }
-        Shape expectedShape = layer->result->output->shape;
-        data.resize(expectedShape.size(), 0);
-        layer->result->output.emplace(Blob::constBlob(expectedShape, data.data()));
+        data.resize(expected_shape.size(), 0);
+        layer->result->output.emplace(Blob::constBlob(expected_shape, data.data()));
     }
 }
 
diff --git a/server/api/GraphBuilder.h b/server/api/GraphBuilder.h
index 41c02327..cbdca510 100644
--- a/server/api/GraphBuilder.h
+++ b/server/api/GraphBuilder.h
@@ -27,14 +27,13 @@ class Graph {
 public:
     Graph() = default;
     void Initialize(crow::json::rvalue modelJson,
-          const std::vector<std::vector<float>>& data,
-          RandomObject* randomInit,
-          OptimizerBase& SGD);
+                    RandomObject* randomInit,
+                    OptimizerBase& SGD,
+                    size_t batch_size);
     ~Graph();
 
-    void OverviewLayers(const crow::json::rvalue& layers, const std::vector<std::vector<float>>& data,
-                        std::unordered_map<int, crow::json::rvalue>& layer_dicts,
-                        std::unordered_map<int, std::vector<float>>& data_dicts);
+    void OverviewLayers(const crow::json::rvalue& layers,
+                        std::unordered_map<int, crow::json::rvalue>& layer_dicts);
 
     void GetEdges(const crow::json::rvalue& connections,
                   std::unordered_map<int, std::vector<int>>& straightEdges,
@@ -45,7 +44,7 @@ class Graph {
                       std::unordered_set<int>& entryNodes,
                       std::vector<int>& layersOrder);
 
-    void ChangeInputData(std::vector<float> data);
+    void ChangeLayersData(std::vector<float> data, BaseLayerType type);
     
     std::vector<Layer*> getLayers(BaseLayerType type) const;
 
diff --git a/server/api/Parser.cpp b/server/api/Parser.cpp
index e1496665..4462f529 100644
--- a/server/api/Parser.cpp
+++ b/server/api/Parser.cpp
@@ -6,15 +6,14 @@ void CHECK_HAS_FIELD(const crow::json::rvalue& layer, const std::string& field)
     }
 }
 
-void ParseCsvData(const std::vector<std::vector<float>>& data, std::vector<float>& instances, std::vector<float>& answers) {
-    // Think about optimization via reservation
-    for (auto& instance : data) {
-        answers.push_back(instance.back());
-        for (int i = 0; i < instance.size() - 1; ++i) {
-            instances.push_back(instance[i]);
-        }
-    }
-}
+// void ParseCsvData(const std::vector<std::vector<float>>& data, std::vector<float>& instances, std::vector<float>& answers) {
+//     instances.reserve(data.size());
+//     answers.reserve(data.size());
+//     for (auto& instance : data) {
+//         answers.push_back(instance.back());
+//         instances.emplace_back(instance.begin(), std::prev(instance.end()));
+//     }
+// }
 
 LinearLayerParameters ParseLinear(const crow::json::rvalue& parameters) {
     size_t inFeatures, outFeatures;
@@ -26,18 +25,23 @@ LinearLayerParameters ParseLinear(const crow::json::rvalue& parameters) {
     inFeatures = static_cast<size_t>(parameters["inFeatures"].i());
     outFeatures = static_cast<size_t>(parameters["outFeatures"].i());
 
-    if (parameters.has("bias")) {
-        bias = parameters["bias"].b();
-    }
+    // if (parameters.has("bias")) {
+    //     std::cout << parameters["bias"] << std::endl;
+    //     bias = parameters["bias"].b();
+    //     std::cout << "done" << std::endl;
+    // }
 
     return LinearLayerParameters{inFeatures, outFeatures, bias};
 }
 
-Data2dLayerParameters ParseData2d(const crow::json::rvalue& parameters) {
-    size_t width;
+Shape ParseData(const crow::json::rvalue& parameters) {
+    std::vector<size_t> shape;
+    shape.reserve(3);
 
-    CHECK_HAS_FIELD(parameters, "width");
+    CHECK_HAS_FIELD(parameters, "shape");
 
-    width = static_cast<size_t>(parameters["width"].i());
-    return Data2dLayerParameters{.width = width};
+    for (auto dim : parameters["shape"]) {
+        shape.push_back(dim.i());
+    }
+    return Shape{std::move(shape)};
 }
diff --git a/server/api/Parser.h b/server/api/Parser.h
index 23662921..a742689c 100644
--- a/server/api/Parser.h
+++ b/server/api/Parser.h
@@ -18,7 +18,7 @@
 
 void CHECK_HAS_FIELD(const crow::json::rvalue& layer, const std::string& field);
 
-void ParseCsvData(const std::vector<std::vector<float>>& data, std::vector<float>& instances, std::vector<float>& answers);
+// void ParseCsvData(const std::vector<std::vector<float>>& data, std::vector<float>& instances, std::vector<float>& answers);
 
 LinearLayerParameters ParseLinear(const crow::json::rvalue& parameters);
-Data2dLayerParameters ParseData2d(const crow::json::rvalue& parameters);
+Shape ParseData(const crow::json::rvalue& parameters);
diff --git a/server/api/server.cpp b/server/api/server.cpp
index 5950b7a3..25b69d00 100644
--- a/server/api/server.cpp
+++ b/server/api/server.cpp
@@ -3,52 +3,63 @@
 #include <vector>
 #include <string>
 #include <sstream>
-
-#include <boost/asio.hpp>
+#include <zip.h>
+#include <filesystem>
+#include <crow_all.h>
 
 #include <cpprest/http_client.h>
 #include <cpprest/filestream.h>
 #include <cpprest/uri.h>
 #include <cpprest/json.h>
 
-#include <crow_all.h>
+#pragma push_macro("U")
+#undef U
+#include "ImageLoader.h"
+#pragma pop_macro("U")
+
 #include "GraphBuilder.h"
 #include "CsvLoader.h"
+#include "DataMarker.h"
 
 using namespace std;
 using namespace crow;
 
 std::string getDataPath(int id) {
-    return "./model_data/data/" + std::to_string(id) + ".csv";
+    return "./model_data/data/" + std::to_string(id);
 }
 
 std::string getPredictPath(int id) {
-    return "./model_data/predict/" + std::to_string(id) + ".csv";
+    return "./model_data/predict/" + std::to_string(id);
 }
 
-web::json::value GetLogs(const std::optional<Blob>& node) {
-    assert(node.has_value() && node.value().shape.dimsCount <= 2);
-    std::vector<web::json::value> values;
-    values.reserve(node.value().shape.size());
-    for (size_t sample_index = 0; sample_index < node.value().shape.rows(); ++sample_index) {
-        for (size_t feature_index = 0; feature_index < node.value().shape.cols(); ++feature_index) {
-            values.push_back(web::json::value::number(node.value()(0, 0, sample_index, feature_index)));
+void GetLogs(const Blob& node, std::vector<web::json::value>& values) {
+    assert(node.shape.dimsCount <= 2);
+    values.reserve(values.size() + node.shape.size());
+    for (size_t sample_index = 0; sample_index < node.shape.rows(); ++sample_index) {
+        for (size_t feature_index = 0; feature_index < node.shape.cols(); ++feature_index) {
+            values.push_back(web::json::value::number(node(0, 0, sample_index, feature_index)));
         }
     }
-    return web::json::value::array(values);
 }
 
-void train(json::rvalue& json, Graph** graph, int model_id, int user_id) {
+void train(json::rvalue& json, Graph** graph, int model_id, int user_id, FileExtension extension) {
     RandomObject initObject(0, 1, 42);
     OptimizerBase SGD = OptimizerBase(0.1);
-    std::vector<std::vector<float>> data = CsvLoader::load_csv(getDataPath(model_id));
+
+    // Should be adopted for DataLoader possibilities
+    std::string path = getDataPath(model_id);
+    if (extension == FileExtension::Csv) {
+        path += "/1.csv";
+    }
+    DataMarker dataMarker = DataMarker(path, extension, 100, 4);
+    DataLoader dataLoader = dataMarker.get_train_loader();
+
+    size_t batch_size = 4;  // hard-coded
     *graph = new Graph();
-    (*graph)->Initialize(json, data, &initObject, SGD);
+    (*graph)->Initialize(json, &initObject, SGD, batch_size);
     std::cout << "Graph is ready!" << std::endl;
 
     auto& lastTrainNode = (*graph)->getLayers(BaseLayerType::TrainOut)[0]->result.value();
-    auto& lastPredictNode = (*graph)->getLayers(BaseLayerType::PredictOut)[0]->result.value().output;
-    auto& targetsNode = (*graph)->getLayers(BaseLayerType::Targets)[0]->result.value().output;
 
     lastTrainNode.forward();
     lastTrainNode.gradient = Blob::ones({{1}});
@@ -57,61 +68,95 @@ void train(json::rvalue& json, Graph** graph, int model_id, int user_id) {
     lastTrainNode.clear();
     Allocator::endVirtualMode();
 
-    size_t buffer_size = 5;
+    size_t buffer_size = 5, actual_size = 0;
+    web::json::value request;
+    request[U("rewrite")] = web::json::value::boolean(true);
     std::vector<web::json::value> targets, outputs;
-    targets.reserve(buffer_size);
-    outputs.reserve(buffer_size);
 
-    size_t max_epochs = 1000;
-    for (size_t epoch = 0; epoch < max_epochs; ++epoch) {
-        auto& result = lastTrainNode.forward();
-        // printf("%ld: %f\n", epoch, result[0][0]);
+    size_t max_epochs = 5;
+    std::pair<std::vector<float>, std::vector<float>> batch;
 
-        outputs.push_back(GetLogs(lastPredictNode));
-        targets.push_back(GetLogs(targetsNode));
+    web::http::client::http_client client(U("http://localhost:3000"));
+    std::ostringstream request_url;
+    request_url << "/update_metrics/" << user_id << "/" << model_id;
 
-        if ((epoch == max_epochs - 1 && outputs.size() > 0) ||
-            outputs.size() == buffer_size) {
+    auto& lastPredictNode = (*graph)->getLayers(BaseLayerType::PredictOut)[0]->result.value().output;
+    auto& targetsNode = (*graph)->getLayers(BaseLayerType::Targets)[0]->result.value().output;
+    for (size_t epoch = 0; epoch < max_epochs; ++epoch) {
+        for (size_t batch_index = 0; batch_index < dataLoader.batch_count(); ++batch_index) {
+            batch = dataLoader.get_raw(batch_index);
+            (*graph)->ChangeLayersData(batch.first, BaseLayerType::Data);
+            (*graph)->ChangeLayersData(batch.second, BaseLayerType::Targets);
 
-            web::json::value json;
-            json["targets"] = web::json::value::array(targets);
-            json["outputs"] = web::json::value::array(outputs);
-            json["label"] = web::json::value::string("train");
-            if (epoch < buffer_size) {
-                json["rewrite"] = web::json::value::boolean(true);
-            }
-            targets.clear();
-            outputs.clear();
+            lastTrainNode.forward();
 
-            std::ostringstream request_url;
-            request_url << "/update_metrics/" << user_id << "/" << model_id;
+            GetLogs(lastPredictNode.value(), targets);
+            GetLogs(targetsNode.value(), outputs);
 
-            web::http::client::http_client client(U("http://localhost:3000"));
-            client.request(web::http::methods::PUT, U(request_url.str()), json);
+            lastTrainNode.gradient = Blob::ones({{1}});
+            lastTrainNode.backward();
+            SGD.step();
+            Allocator::endSession();
+            lastTrainNode.clear();
         }
 
-        // lastTrainNode.gradient = result;
-        lastTrainNode.gradient = Blob::ones({{1}});
-        lastTrainNode.backward();
-        SGD.step();
-        lastTrainNode.clear();
+        request[U("targets")][actual_size] = web::json::value::array(targets);
+        request[U("outputs")][actual_size] = web::json::value::array(outputs);
+        targets.clear();
+        outputs.clear();
+        ++actual_size;
+
+        if ((epoch == max_epochs - 1 && actual_size > 0) ||
+            actual_size == buffer_size) {
+            request[U("label")] = web::json::value::string("train");
+            client.request(web::http::methods::PUT, U(request_url.str()), request);
+            request = web::json::value();
+            actual_size = 0;
+        }
     }
 }
 
-void predict(int model_id, Graph* graph, std::vector<float>& answer) {
-    std::vector<std::vector<float>> predict_data = CsvLoader::load_csv(getPredictPath(model_id));
-    graph->ChangeInputData(predict_data[0]);
+void predict(int model_id, Graph* graph, std::vector<float>& answer, FileExtension extension) {
+    std::vector<std::vector<float>> predict_data;
+    if (extension == FileExtension::Csv) {
+        predict_data = CsvLoader::load_csv(getPredictPath(model_id) + "/1.csv");
+    }
+    else {
+        predict_data = {ImageLoader::load_image((getPredictPath(model_id) + "/1.png").c_str())};
+    }
+    graph->ChangeLayersData(predict_data[0], BaseLayerType::Data);
 
     // Пока не думаем о нескольких выходах (!) Hard-coded
     auto& lastNode = graph->getLayers(BaseLayerType::PredictOut)[0]->result.value();
     lastNode.clear();
-    const Blob& result = lastNode.forward();
+    lastNode.forward();
+
+    auto& result = lastNode.output.value();
+    assert(result.shape.dimsCount == 2);
 
-    answer.reserve(result.shape.rows() * result.shape.cols());
-    for (size_t j = 0; j < result.shape.rows(); ++j) {
-        for (size_t i = 0; i < result.shape.cols(); ++i) {
-            answer.push_back(result(j, i));
-            std::cout << result(j, i) << std::endl;
+    answer.reserve(result.shape.cols());
+    for (size_t i = 0; i < result.shape.cols(); ++i) {
+        answer.push_back(result(0, i));
+        std::cout << result(0, i) << std::endl;
+    }
+}
+
+void extract_from_zip(std::string path, std::string root) {
+    zip_t* z;
+    int err;
+    z = zip_open(path.c_str(), 0, &err);
+    if (z == nullptr) {
+        throw std::runtime_error("File doesn't exist");
+    }
+    zip_stat_t info;
+    for (int i = 0; i < zip_get_num_entries(z, ZIP_FL_UNCHANGED); ++i) {
+        if (zip_stat_index(z, i, 0, &info) == 0) {
+            ofstream fout(root + "/" + info.name, ios::binary);
+            zip_file* file = zip_fopen_index(z, i, 0);
+            char *file_data = new char[info.size];
+            zip_fread(file, file_data, info.size);
+            fout.write(file_data, info.size);
+            delete[] file_data;
         }
     }
 }
@@ -126,20 +171,20 @@ int main(int argc, char *argv[]) {
     SimpleApp app;
 
     std::map<int, Graph*> sessions;
+    std::map<int, FileExtension> file_types;
 
-    CROW_ROUTE(app, "/predict/<int>").methods(HTTPMethod::POST)
+    CROW_ROUTE(app, "/predict/<int>").methods(HTTPMethod::GET)
     ([&](const request& req, int model_id) -> response {
+        (void)req;
         if (sessions.find(model_id) == sessions.end()) return response(status::METHOD_NOT_ALLOWED, "Not trained");
         std::vector<float> answer;
         try {
-            predict(model_id, sessions[model_id], answer);
+            predict(model_id, sessions[model_id], answer, file_types[model_id]);
         } catch (const std::runtime_error &err) {
+            std::cout << err.what() << std::endl;
             return response(status::BAD_REQUEST, "Invalid body");
         }
-        json::wvalue response;
-        for (int i = 0; i < answer.size(); ++i) {
-            response[i] = answer[i];
-        }
+        json::wvalue response = answer[0];
         return crow::response(status::OK, response);
     });
 
@@ -154,32 +199,59 @@ int main(int argc, char *argv[]) {
         }
         Graph* g = nullptr;
 
-        train(body, &g, user_id, model_id);
+        train(body, &g, model_id, user_id, file_types[model_id]);
         sessions[model_id] = g;
         return response(status::OK, "done");
     });
 
-    //curl -X POST -F "InputFile=@filename" http://0.0.0.0:2000/upload_data/1/0 (last can be 1)
+    // curl -X POST -F "InputFile=@filename" http://0.0.0.0:2000/upload_data/1/0/0
+    // Second argument is for request type (train or predict), third - for file type (csv or zip)
     CROW_ROUTE(app, "/upload_data/<int>/<int>").methods(HTTPMethod::Post)
     ([&](const request& req, int model_id, int type) -> response {
-        crow::multipart::message file_message(req);
-        std::string path;
+        string content_type = req.get_header_value("Content-Type");
+        if (content_type.empty()) return response(status::BAD_REQUEST, "No Content-Type header provided");
+
+        std::string path, root;
         if (type == 0) {
             path = getDataPath(model_id);
         }
         else {
             path = getPredictPath(model_id);
         }
-        std::ofstream out_file(path);
-        if (!out_file) {
-            return response(status::INTERNAL_SERVER_ERROR, "Failed to open file for storage");
+        root = path;
+
+        if (std::filesystem::exists(path)) {
+            std::filesystem::remove_all(path);
         }
-        auto content = file_message.part_map.find("InputFile");
-        if (content == file_message.part_map.end()) {
-            return response(status::BAD_REQUEST, "No file provided");
+        std::filesystem::create_directory(path);
+
+        if (content_type == "text/csv") {
+            path += "/1.csv";
+            file_types[model_id] = FileExtension::Csv;
         }
-        out_file << (*content).second.body;
+        else if (content_type == "application/zip") {
+            path += "/1.zip";
+            file_types[model_id] = FileExtension::Png;
+        } else if (content_type == "image/png") {
+            path += "/1.png";
+            file_types[model_id] = FileExtension::Png;
+        } else {
+            return response(status::BAD_REQUEST, "Invalid content type: " + content_type);
+        }
+
+        std::ofstream out_file(path);
+        out_file << req.body;
         out_file.close();
+
+        if (content_type == "application/zip" && type == 0) {
+            try {
+                extract_from_zip(path, root);
+            }
+            catch (...) {
+                return response(status::INTERNAL_SERVER_ERROR, "Error in extracting from zip");
+            }
+        }
+
         return crow::response(status::OK, "done");
     });
 
diff --git a/server/core/Allocator.cpp b/server/core/Allocator.cpp
index 2ae4a83c..faa6fa29 100644
--- a/server/core/Allocator.cpp
+++ b/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);
diff --git a/server/core/Layer.cpp b/server/core/Layer.cpp
index 7f143a42..708f3526 100644
--- a/server/core/Layer.cpp
+++ b/server/core/Layer.cpp
@@ -32,9 +32,14 @@ ReLULayer::ReLULayer(const std::vector<TensorRef>& args) {
     result = Tensor(relu, {args[0]});
 }
 
-Data2dLayer::Data2dLayer(const Data2dLayerParameters& params, const std::vector<float>& values)
-    : width(params.width) {
-    result = Tensor(Blob::constBlob({{params.height, width}}, values.data()));
+DataLayer::DataLayer(const Shape& shape, const std::vector<float>& values) {
+    result = Tensor(Blob::constBlob(shape, values.data()));
+}
+
+DataLayer::DataLayer(const Shape& shape, size_t batch_size) {
+    std::vector<size_t> dims = shape.getDims();
+    dims[0] = batch_size;
+    result = Tensor(Blob::constRandomBlob(Shape(std::move(dims)), nullptr));
 }
 
 MSELoss::MSELoss(const std::vector<TensorRef>& args) : mean({0, 1, 2, 3}) {
diff --git a/server/core/Layer.h b/server/core/Layer.h
index 8d9cbb50..bdca7c95 100644
--- a/server/core/Layer.h
+++ b/server/core/Layer.h
@@ -5,6 +5,7 @@
 
 #include "RandomInit.h"
 #include "Tensor.h"
+#include "Shape.h"
 #include "Parameters.h"
 
 
@@ -16,10 +17,10 @@ class Layer {
     std::vector<TensorRef> layerOperationParams;
 };
 
-class Data2dLayer: public Layer {
+class DataLayer: public Layer {
 public:
-    size_t width;
-    Data2dLayer(const Data2dLayerParameters& params, const std::vector<float>& values);
+    DataLayer(const Shape& params, const std::vector<float>& values);
+    DataLayer(const Shape& params, size_t batch_size);
 };
 
 class LinearLayer: public Layer {
diff --git a/server/core/LazyBlob.cpp b/server/core/LazyBlob.cpp
index 3b6cd432..601cf675 100644
--- a/server/core/LazyBlob.cpp
+++ b/server/core/LazyBlob.cpp
@@ -388,7 +388,7 @@ const LazyBlob& operator / (const LazyBlob &a, float b) {
 Blob& operator += (Blob& a, const LazyBlob& b) {
     assert(a.shape == b.shape());
 
-    std::vector<size_t> size =  a.shape.getDims();
+    std::vector<size_t> size = a.shape.getDims();
     map(size, [&](size_t k, size_t l, size_t i, size_t j) {
        a(k, l, i, j) += b(k, l, i, j);
     });
@@ -398,7 +398,7 @@ Blob& operator += (Blob& a, const LazyBlob& b) {
 Blob& operator -= (Blob& a, const LazyBlob& b) {
     assert(a.shape == b.shape());
 
-    std::vector<size_t> size =  a.shape.getDims();
+    std::vector<size_t> size = a.shape.getDims();
     map(size, [&](size_t k, size_t l, size_t i, size_t j) {
        a(k, l, i, j) -= b(k, l, i, j);
     });
@@ -408,7 +408,7 @@ Blob& operator -= (Blob& a, const LazyBlob& b) {
 Blob& operator *= (Blob& a, const LazyBlob& b) {
     assert(a.shape == b.shape());
 
-    std::vector<size_t> size =  a.shape.getDims();
+    std::vector<size_t> size = a.shape.getDims();
     map(size, [&](size_t k, size_t l, size_t i, size_t j) {
        a(k, l, i, j) *= b(k, l, i, j);
     });
diff --git a/server/core/Optimizer.cpp b/server/core/Optimizer.cpp
index 213f5745..7c57eafe 100644
--- a/server/core/Optimizer.cpp
+++ b/server/core/Optimizer.cpp
@@ -5,13 +5,18 @@
 #include "Tensor.h"
 #include "Optimizer.h"
 
+OptimizerBase::OptimizerBase(float lr) : lr(lr) {
+}
+
 void OptimizerBase::append(std::vector<TensorRef>& newParams) {
-    params.reserve(params.size() + newParams.size() * sizeof(TensorRef));
-    params.insert(params.end(),newParams.begin(),newParams.end());
+    params.reserve(params.size() + newParams.size());
+    params.insert(params.end(), newParams.begin(), newParams.end());
 }
 
 void OptimizerBase::step() {
     for (int i = 0; i < params.size(); i++) {
-        *params[i].get().output -= lr * *params[i].get().gradient;
+        if (params[i].get().gradient.has_value() && params[i].get().output.has_value()) {
+            params[i].get().output.value() -= lr * params[i].get().gradient.value();
+        }
     }
 }
\ No newline at end of file
diff --git a/server/core/Optimizer.h b/server/core/Optimizer.h
index f5cb0e75..50843a79 100644
--- a/server/core/Optimizer.h
+++ b/server/core/Optimizer.h
@@ -5,12 +5,12 @@
 #include "Tensor.h"
 
 class OptimizerBase {
-
-const float lr;
-std::vector<TensorRef> params;
-
 public:
-    OptimizerBase(float lr) : lr(lr), params() {};
+    OptimizerBase(float lr);
     void append(std::vector<TensorRef>& newParams);
     void step();
-};
\ No newline at end of file
+
+private:
+    const float lr;
+    std::vector<TensorRef> params;
+};
diff --git a/server/core/Parameters.h b/server/core/Parameters.h
index bb20cc20..27fce33a 100644
--- a/server/core/Parameters.h
+++ b/server/core/Parameters.h
@@ -5,8 +5,3 @@ struct LinearLayerParameters {
     std::size_t outFeatures;
     bool bias;
 };
-
-struct Data2dLayerParameters {
-    std::size_t width;
-    std::size_t height;
-};
diff --git a/server/tests/GraphTests.cpp b/server/tests/GraphTests.cpp
index 4e5f242d..a038ab0d 100644
--- a/server/tests/GraphTests.cpp
+++ b/server/tests/GraphTests.cpp
@@ -38,19 +38,13 @@ void CheckFullProcess(const std::string& json_path,
     //                     PARSE LAYERS IDS
     // #################################################################
     std::unordered_map<int, crow::json::rvalue> layerDicts;
-    std::unordered_map<int, std::vector<float>> dataDicts;
-    g.OverviewLayers(layersJson, data, layerDicts, dataDicts);
+    g.OverviewLayers(layersJson, layerDicts);
 
     CHECK(layerDicts.size() == layer_ids_expected.size());
     for (auto id : layer_ids_expected) {
         CHECK(layerDicts.find(id) != layerDicts.end());
     }
 
-    CHECK(dataDicts.size() == data_ids_expected.size());
-    for (auto id : data_ids_expected) {
-        CHECK(dataDicts.find(id) != dataDicts.end());
-    }
-
     // #################################################################
     //                         PARSE EDGES
     // #################################################################
diff --git a/server/tests/linear_relu_linear_mse.json b/server/tests/linear_relu_linear_mse.json
index f1d872ac..f89468be 100644
--- a/server/tests/linear_relu_linear_mse.json
+++ b/server/tests/linear_relu_linear_mse.json
@@ -31,7 +31,7 @@
                 "id": 0,
                 "type": "Data",
                 "parameters": {
-                    "width": 2
+                    "shape": [2]
                 }
             },
             {
@@ -60,7 +60,7 @@
                 "id": 4,
                 "type": "Target",
                 "parameters": {
-                    "width": 1
+                    "shape": [1]
                 }
             },
             {