Added model loading in train_cli.py

Fixed model export
Set kernels for alpha-vile-large

DeepCrazyhouse/src/domain/neural_net/architectures/pytorch/alpha_vile.py

@@ -46,9 +46,29 @@ def get_alpha_vile_model(args, model_size='normal'):
 
     kernels = [3] * depth
     end_idx = int(len(kernels) * kernel_5_ratio + 0.5)
-    for idx in range(end_idx):
-        kernels[idx] = 5
-    random.shuffle(kernels)
+
+    if model_size == 'large':
+        kernels[1] = 5
+        kernels[6] = 5
+        kernels[7] = 5
+        kernels[9] = 5
+        kernels[10] = 5
+        kernels[14] = 5
+        kernels[18] = 5
+        kernels[19] = 5
+        kernels[23] = 5
+        kernels[25] = 5
+        kernels[26] = 5
+        kernels[27] = 5
+        kernels[28] = 5
+        kernels[29] = 5
+        kernels[33] = 5
+        kernels[34] = 5
+        kernels[35] = 5
+    else:
+        for idx in range(end_idx):
+            kernels[idx] = 5
+        random.shuffle(kernels)
 
     use_transformers = [False] * len(kernels)
     if nb_transformers > 0:

DeepCrazyhouse/src/training/train_cli.py

@@ -16,14 +16,15 @@
 import sys
 import torch
 import logging
+from pathlib import Path
 
 sys.path.insert(0, '../../../')
 
 from DeepCrazyhouse.src.runtime.color_logger import enable_color_logging
 from DeepCrazyhouse.configs.train_config import TrainConfig, TrainObjects
 from DeepCrazyhouse.src.training.train_cli_util import create_pytorch_model, get_validation_data, fill_train_objects,\
     print_model_summary, export_best_model_state, fill_train_config, export_configs, create_export_dirs, export_cmd_args
-from DeepCrazyhouse.src.training.trainer_agent_pytorch import TrainerAgentPytorch
+from DeepCrazyhouse.src.training.trainer_agent_pytorch import TrainerAgentPytorch, load_torch_state
 
 
 def parse_args(train_config: TrainConfig):
@@ -91,6 +92,10 @@ def main():
 
     train_objects = TrainObjects()
     fill_train_objects(train_config, train_objects)
+    if train_config.tar_file != "":
+        print("load model weights")
+        load_torch_state(model, torch.optim.SGD(model.parameters(), lr=train_config.max_lr), Path(train_config.tar_file),
+                         train_config.device_id)
 
     create_export_dirs(train_config)
     export_configs(args, train_config)

DeepCrazyhouse/src/training/train_cli_util.py

@@ -38,7 +38,7 @@
     LinearWarmUp, MomentumSchedule
 from DeepCrazyhouse.src.training.train_util import get_metrics
 from DeepCrazyhouse.src.training.trainer_agent_pytorch import save_torch_state, export_to_onnx, get_context,\
-    get_data_loader
+    get_data_loader, load_torch_state
 
 
 class Args:
@@ -232,6 +232,12 @@ def export_best_model_state(k_steps_best: int, k_steps_final: int, model, policy
     shutil.copy(model_tar_path, best_model_tar_path)
 
     # ## Convert to onnx
+    print("load current best model")
+    load_torch_state(model, torch.optim.SGD(model.parameters(), lr=train_config.max_lr), Path(model_tar_path),
+                     train_config.device_id)
+
+    if hasattr(model, "merge_bn"):
+        model.merge_bn()
     convert_model_to_onnx(input_shape, k_steps_best, model, model_name, train_config)
 
     print("Saved weight & onnx files of the best model to %s" % (train_config.export_dir + "best-model"))