model_to_txt.py 5.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143
  1. import os
  2. import keras
  3. import psutil
  4. import configparser
  5. import os
  6. import sys
  7. sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
  8. import scripts.tools.utils as tools_utils
  9. import argparse
  10. import json
  11. class LossHistory(keras.callbacks.Callback):
  12. def on_train_begin(self, logs=None):
  13. self.losses = []
  14. self.acc = []
  15. self.mem_info_list = []
  16. def on_batch_end(self, batch, logs=None):
  17. process = psutil.Process()
  18. self.losses.append(logs.get('loss'))
  19. self.acc.append(logs.get('acc'))
  20. self.mem_info_list.append(process.memory_info().rss / (1024**3))
  21. def save_log_txt(model, path, name, bk, x_train, y_train):
  22. history_loss = LossHistory()
  23. model.fit(x_train, y_train, epochs=2, batch_size=1024, validation_split=0.2, verbose=1, callbacks=[history_loss])
  24. model_name = name.split("/")[-1].split("_")[0]
  25. method_name = name.split("/")[-1].split("_")[1]
  26. valid_acc = [acc for acc in history_loss.acc if acc is not None]
  27. result_data = {
  28. "model": model_name,
  29. "method": method_name,
  30. "result": {
  31. "Losses": sum(history_loss.losses)/len(history_loss.losses),
  32. "Accuracy": sum(valid_acc)/len(valid_acc),
  33. "MemoryInfoList": sum(history_loss.mem_info_list)/len(history_loss.mem_info_list)
  34. }
  35. }
  36. if os.path.exists(path):
  37. with open(path, 'r') as json_file:
  38. data = json.load(json_file)
  39. else:
  40. data = []
  41. data.append(result_data)
  42. with open(path, 'w') as json_file:
  43. json.dump(data, json_file, indent=4)
  44. def custom_objects():
  45. def no_activation(x):
  46. return x
  47. def leakyrelu(x):
  48. import keras.backend as K
  49. return K.relu(x, alpha=0.01)
  50. objects = {}
  51. objects['no_activation'] = no_activation
  52. objects['leakyrelu'] = leakyrelu
  53. return objects
  54. def model_to_txt1(model_path, bk):
  55. cur_model = keras.models.load_model(model_path, custom_objects=custom_objects())
  56. cur_model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
  57. data = tools_utils.DataUtils
  58. if 'svhn' in model_path or 'fashion2' in model_path:
  59. txt_path = model_path.split("\\")[-1][:-5].split("mut_model")[0]+bk+".json"
  60. else:
  61. txt_path = model_path.split("\\")[-1][:-3].split("mut_model")[0]+bk+".json"
  62. if 'svhn' in model_path or 'fashion2' in model_path:
  63. model_path = model_path.split("\\")[-1][:-5]
  64. else:
  65. model_path = model_path.split("\\")[-1][:-3]
  66. data_path = model_path.split("/")[-1]
  67. x_test, y_test = data.get_data_by_exp(data_path)
  68. save_log_txt(cur_model,txt_path,model_path,bk,x_test, y_test)
  69. if __name__ == "__main__":
  70. """Parser of command args"""
  71. parse = argparse.ArgumentParser()
  72. parse.add_argument("--backend", type=str, help="name of backends")
  73. parse.add_argument("--model_path", type=str, help="redis db port")
  74. parse.add_argument("--root_dir", type=str, help="redis db port")
  75. flags, unparsed = parse.parse_known_args(sys.argv[1:])
  76. """Load Configuration"""
  77. lemon_cfg = configparser.ConfigParser()
  78. # lemon_cfg.read(f"./config/{flags.config_name}")
  79. conf_path = os.path.join(os.path.dirname(os.getcwd()), "config", "demo.conf")
  80. lemon_cfg.read(conf_path)
  81. parameters = lemon_cfg['parameters']
  82. gpu_ids = parameters['gpu_ids']
  83. gpu_list = parameters['gpu_ids'].split(",")
  84. """Init cuda"""
  85. os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
  86. os.environ["CUDA_VISIBLE_DEVICES"] = gpu_ids
  87. """Switch backend"""
  88. bk_list = ['tensorflow', 'theano', 'cntk', 'mxnet']
  89. bk = flags.backend
  90. print('.........................',type(bk))
  91. os.environ['KERAS_BACKEND'] = bk
  92. os.environ['PYTHONHASHSEED'] = '0'
  93. if bk == 'tensorflow':
  94. os.environ["TF_CPP_MIN_LOG_LEVEL"] = '2' # 只显示 warning 和 Error
  95. import tensorflow as tf
  96. if bk == 'theano':
  97. if len(gpu_list) == 2:
  98. os.environ[
  99. 'THEANO_FLAGS'] = f"device=cuda,contexts=dev{gpu_list[0]}->cuda{gpu_list[0]};dev{gpu_list[1]}->cuda{gpu_list[1]}," \
  100. f"force_device=True,floatX=float32,lib.cnmem=1"
  101. else:
  102. os.environ['THEANO_FLAGS'] = f"device=cuda,contexts=dev{gpu_list[0]}->cuda{gpu_list[0]}," \
  103. f"force_device=True,floatX=float32,lib.cnmem=1"
  104. batch_size = 32
  105. import theano as th
  106. mylogger.info(th.__version__)
  107. if bk == "cntk":
  108. batch_size = 32
  109. from cntk.device import try_set_default_device, gpu
  110. try_set_default_device(gpu(int(gpu_list[0])))
  111. import cntk as ck
  112. if bk == "mxnet":
  113. batch_size = 32
  114. import mxnet as mxnet
  115. from keras import backend as K
  116. import keras
  117. print("Using {} as backend for states extraction| {} is wanted".format(K.backend(),bk))
  118. """Get model hidden output on selected_index data on specific backend"""
  119. model_to_txt1(flags.model_path, bk)