Upd new model code

VSharp.ML.AIAgent/common/constants.py

@@ -11,12 +11,8 @@ def _build_bar_format() -> str:
     return f"{custom_left} {custom_bar} - {custom_info}"
 
 
-IMPORTED_FULL_MODEL_PATH = Path(
-    "ml/imported/GNN_state_pred_het_full_TAGConv_20e_2xAll_10h"
-)
-IMPORTED_DICT_MODEL_PATH = Path(
-    "ml/imported/GNN_state_pred_het_dict_TAGConv_20e_2xAll_10h"
-)
+IMPORTED_FULL_MODEL_PATH = Path("ml/imported/GNN_state_pred_het_full_compact.zip")
+IMPORTED_DICT_MODEL_PATH = Path("ml/imported/GNN_state_pred_het_dict_compact.zip")
 
 BASE_REPORT_DIR = Path("./report")
 TABLES_LOG_FILE = BASE_REPORT_DIR / "tables.log"

VSharp.ML.AIAgent/config.py

@@ -6,6 +6,7 @@
 import torch
 
 import ml.model_modified
+import ml.models
 
 
 class GeneralConfig:
@@ -17,6 +18,9 @@ class GeneralConfig:
     MAX_STEPS = 100
     MUTATION_PERCENT_GENES = 4
     LOGGER_LEVEL = logging.INFO
+    IMPORT_MODEL_INIT = lambda: ml.models.StateModelEncoder(
+        hidden_channels=64, out_channels=8
+    )
     MODEL_INIT = lambda: ml.model_modified.StateModelEncoderExport(
         hidden_channels=64, out_channels=8
     )

VSharp.ML.AIAgent/connection/broker_conn/requests.py

@@ -10,7 +10,7 @@
 
 def acquire_instance() -> ServerInstanceInfo:
     response, content = httplib2.Http().request(WebsocketSourceLinks.GET_WS)
-    aquired_instance = ServerInstanceInfo.from_json(content.decode("utf-8"))
+    aquired_instance = ServerInstanceInfo.from_json(json.loads(content.decode("utf-8")))
     logging.info(f"acquired ws: {aquired_instance}")
     return aquired_instance
 

VSharp.ML.AIAgent/main.py

@@ -1,11 +1,13 @@
 import learning.entry_point as ga
 import ml.onnx.onnx_import
+from common.constants import IMPORTED_DICT_MODEL_PATH
 from config import GeneralConfig
 from learning.selection.crossover_type import CrossoverType
 from learning.selection.mutation_type import MutationType
 from learning.selection.parent_selection_type import ParentSelectionType
 from ml.utils import (
     create_population,
+    load_model,
     model_weights_with_last_layer,
     model_weights_with_random_last_layer,
 )
@@ -16,9 +18,10 @@ def main():
 
     model.forward(*ml.onnx.onnx_import.create_torch_dummy_input())
 
-    random_population = create_population(lo=-5, hi=5, model=model, population_size=4)
+    random_population = create_population(lo=-5, hi=5, model=model, population_size=60)
     with_random_last_layer = [
-        model_weights_with_random_last_layer(lo=-1, hi=1, model=model) for _ in range(2)
+        model_weights_with_random_last_layer(lo=-1, hi=1, model=model)
+        for _ in range(18)
     ]
     with_last_layer1 = model_weights_with_last_layer(
         [
@@ -31,7 +34,7 @@ def main():
             0.95478059636744,
             0.27937866719070503,
         ],
-        model,
+        load_model(path=IMPORTED_DICT_MODEL_PATH, model=GeneralConfig.MODEL_INIT()),
     )
     with_last_layer2 = model_weights_with_last_layer(
         [
@@ -44,7 +47,7 @@ def main():
             0.9555442824877577,
             0.2793786892860371,
         ],
-        model,
+        load_model(path=IMPORTED_DICT_MODEL_PATH, model=GeneralConfig.MODEL_INIT()),
     )
 
     initial_population = [

VSharp.ML.AIAgent/ml/utils.py

@@ -1,5 +1,6 @@
 import copy
 import random
+from pathlib import Path
 
 import numpy
 import pygad.torchga
@@ -14,6 +15,12 @@
 from ml.onnx.onnx_import import create_torch_dummy_input
 
 
+def load_model(path: Path, model: torch.nn.Module):
+    model.load_state_dict(torch.load(path))
+    model.eval()
+    return model
+
+
 def model_weights_with_last_layer(
     last_layer_weights: list[float], model: ml.models.StateGNNEncoderConvEdgeAttr
 ) -> npt.NDArray: