From c1072979563058d560abe17951e680291d8f6cc0 Mon Sep 17 00:00:00 2001
From: luopengting <luopengting@huawei.com>
Date: Wed, 22 Apr 2020 17:32:27 +0800
Subject: [PATCH] add device_num

---
 mindinsight/lineagemgr/querier/query_model.py  | 1 +
 tests/st/func/lineagemgr/api/test_model_api.py | 3 +++
 tests/ut/lineagemgr/querier/test_querier.py    | 3 +++
 3 files changed, 7 insertions(+)

diff --git a/mindinsight/lineagemgr/querier/query_model.py b/mindinsight/lineagemgr/querier/query_model.py
index b797cdc..3823cd7 100644
--- a/mindinsight/lineagemgr/querier/query_model.py
+++ b/mindinsight/lineagemgr/querier/query_model.py
@@ -35,6 +35,7 @@ FIELD_MAPPING = {
     "learning_rate": Field('hyper_parameters', 'learning_rate'),
     "epoch": Field('hyper_parameters', 'epoch'),
     "batch_size": Field('hyper_parameters', 'batch_size'),
+    "device_num": Field('hyper_parameters', 'device_num'),
     "loss": Field('algorithm', 'loss'),
     "model_size": Field('model', 'size'),
     "dataset_mark": Field('dataset_mark', None),
diff --git a/tests/st/func/lineagemgr/api/test_model_api.py b/tests/st/func/lineagemgr/api/test_model_api.py
index c824ee6..d50cad3 100644
--- a/tests/st/func/lineagemgr/api/test_model_api.py
+++ b/tests/st/func/lineagemgr/api/test_model_api.py
@@ -76,6 +76,7 @@ LINEAGE_FILTRATION_EXCEPT_RUN = {
         'learning_rate': 0.11999999731779099,
         'epoch': 10,
         'batch_size': 32,
+        'device_num': 2,
         'loss': 0.029999999329447746,
         'model_size': 64,
         'metric': {},
@@ -97,6 +98,7 @@ LINEAGE_FILTRATION_RUN1 = {
         'learning_rate': 0.11999999731779099,
         'epoch': 14,
         'batch_size': 32,
+        'device_num': 2,
         'loss': None,
         'model_size': 64,
         'metric': {
@@ -120,6 +122,7 @@ LINEAGE_FILTRATION_RUN2 = {
         'learning_rate': None,
         'epoch': None,
         'batch_size': None,
+        'device_num': None,
         'loss': None,
         'model_size': None,
         'metric': {
diff --git a/tests/ut/lineagemgr/querier/test_querier.py b/tests/ut/lineagemgr/querier/test_querier.py
index 6059b24..c1562bb 100644
--- a/tests/ut/lineagemgr/querier/test_querier.py
+++ b/tests/ut/lineagemgr/querier/test_querier.py
@@ -93,6 +93,7 @@ def create_filtration_result(summary_dir, train_event_dict,
             "learning_rate": train_event_dict['train_lineage']['hyper_parameters']['learning_rate'],
             "epoch": train_event_dict['train_lineage']['hyper_parameters']['epoch'],
             "batch_size": train_event_dict['train_lineage']['hyper_parameters']['batch_size'],
+            "device_num": train_event_dict['train_lineage']['hyper_parameters']['device_num'],
             "loss": train_event_dict['train_lineage']['algorithm']['loss'],
             "model_size": train_event_dict['train_lineage']['model']['size'],
             "metric": metric_dict,
@@ -208,6 +209,7 @@ LINEAGE_FILTRATION_5 = {
             event_data.EVENT_TRAIN_DICT_5['train_lineage']['hyper_parameters']['learning_rate'],
         "epoch": event_data.EVENT_TRAIN_DICT_5['train_lineage']['hyper_parameters']['epoch'],
         "batch_size": event_data.EVENT_TRAIN_DICT_5['train_lineage']['hyper_parameters']['batch_size'],
+        "device_num": event_data.EVENT_TRAIN_DICT_5['train_lineage']['hyper_parameters']['device_num'],
         "loss": event_data.EVENT_TRAIN_DICT_5['train_lineage']['algorithm']['loss'],
         "model_size": event_data.EVENT_TRAIN_DICT_5['train_lineage']['model']['size'],
         "metric": {},
@@ -231,6 +233,7 @@ LINEAGE_FILTRATION_6 = {
         "learning_rate": None,
         "epoch": None,
         "batch_size": None,
+        "device_num": None,
         "loss": None,
         "model_size": None,
         "metric": event_data.METRIC_5,
-- 
GitLab