model.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298
  1. import time
  2. import os
  3. import shutil
  4. import pickle
  5. from os import path as osp
  6. from .utils import set_folder_status, TaskStatus, copy_pretrained_model, PretrainedModelStatus
  7. from . import workspace_pb2 as w
  8. def list_pretrained_models(workspace):
  9. """列出预训练模型列表
  10. """
  11. pretrained_model_list = list()
  12. for id in workspace.pretrained_models:
  13. pretrained_model = workspace.pretrained_models[id]
  14. model_id = pretrained_model.id
  15. model_name = pretrained_model.name
  16. model_model = pretrained_model.model
  17. model_type = pretrained_model.type
  18. model_pid = pretrained_model.pid
  19. model_tid = pretrained_model.tid
  20. model_create_time = pretrained_model.create_time
  21. model_path = pretrained_model.path
  22. attr = {
  23. 'id': model_id,
  24. 'name': model_name,
  25. 'model': model_model,
  26. 'type': model_type,
  27. 'pid': model_pid,
  28. 'tid': model_tid,
  29. 'create_time': model_create_time,
  30. 'path': model_path
  31. }
  32. pretrained_model_list.append(attr)
  33. return {'status': 1, "pretrained_models": pretrained_model_list}
  34. def create_pretrained_model(data, workspace, monitored_processes):
  35. """根据request创建预训练模型。
  36. Args:
  37. data为dict,key包括
  38. 'pid'所属项目id, 'tid'所属任务id,'name'预训练模型名称,
  39. 'source_path' 原模型路径, 'eval_results'(可选) 评估结果数据
  40. """
  41. time_array = time.localtime(time.time())
  42. create_time = time.strftime("%Y-%m-%d %H:%M:%S", time_array)
  43. id = workspace.max_pretrained_model_id + 1
  44. workspace.max_pretrained_model_id = id
  45. if id < 10000:
  46. id = 'PM%04d' % id
  47. else:
  48. id = 'PM{}'.format(id)
  49. pid = data['pid']
  50. tid = data['tid']
  51. name = data['name']
  52. source_path = data['source_path']
  53. assert pid in workspace.projects, "【预训练模型创建】项目ID'{}'不存在.".format(pid)
  54. assert tid in workspace.tasks, "【预训练模型创建】任务ID'{}'不存在.".format(tid)
  55. assert not id in workspace.pretrained_models, "【预训练模型创建】预训练模型'{}'已经被占用.".format(
  56. id)
  57. assert osp.exists(source_path), "原模型路径不存在: {}".format(source_path)
  58. path = osp.join(workspace.path, 'pretrain', id)
  59. if not osp.exists(path):
  60. os.makedirs(path)
  61. set_folder_status(path, PretrainedModelStatus.XPINIT)
  62. params = {'tid': tid}
  63. from .project.task import get_task_params
  64. ret = get_task_params(params, workspace)
  65. train_params = ret['train']
  66. model_structure = train_params.model
  67. if hasattr(train_params, "backbone"):
  68. model_structure = "{}-{}".format(model_structure,
  69. train_params.backbone)
  70. if hasattr(train_params, "with_fpn"):
  71. if train_params.with_fpn:
  72. model_structure = "{}-{}".format(model_structure, "WITH_FPN")
  73. pm = w.PretrainedModel(
  74. id=id,
  75. name=name,
  76. model=model_structure,
  77. type=workspace.projects[pid].type,
  78. pid=pid,
  79. tid=tid,
  80. create_time=create_time,
  81. path=path)
  82. workspace.pretrained_models[id].CopyFrom(pm)
  83. # 保存评估结果
  84. if 'eval_results' in data:
  85. with open(osp.join(source_path, "eval_res.pkl"), "wb") as f:
  86. pickle.dump(data['eval_results'], f)
  87. # 拷贝训练参数文件
  88. task_path = workspace.tasks[tid].path
  89. task_params_path = osp.join(task_path, 'params.pkl')
  90. if osp.exists(task_params_path):
  91. shutil.copy(task_params_path, path)
  92. # 拷贝数据集信息文件
  93. did = workspace.projects[pid].did
  94. dataset_path = workspace.datasets[did].path
  95. dataset_info_path = osp.join(dataset_path, "statis.pkl")
  96. if osp.exists(dataset_info_path):
  97. # 写入部分数据集信息
  98. with open(dataset_info_path, "rb") as f:
  99. dataset_info_dict = pickle.load(f)
  100. dataset_info_dict['name'] = workspace.datasets[did].name
  101. dataset_info_dict['desc'] = workspace.datasets[did].desc
  102. with open(dataset_info_path, "wb") as f:
  103. pickle.dump(dataset_info_dict, f)
  104. shutil.copy(dataset_info_path, path)
  105. # copy from source_path to path
  106. proc = copy_pretrained_model(source_path, path)
  107. monitored_processes.put(proc.pid)
  108. return {'status': 1, 'pmid': id}
  109. def delete_pretrained_model(data, workspace):
  110. """删除pretrained_model。
  111. Args:
  112. data为dict,
  113. key包括'pmid'预训练模型id
  114. """
  115. pmid = data['pmid']
  116. assert pmid in workspace.pretrained_models, "预训练模型ID'{}'不存在.".format(pmid)
  117. if osp.exists(workspace.pretrained_models[pmid].path):
  118. shutil.rmtree(workspace.pretrained_models[pmid].path)
  119. del workspace.pretrained_models[pmid]
  120. return {'status': 1}
  121. def create_exported_model(data, workspace):
  122. """根据request创建已发布模型。
  123. Args:
  124. data为dict,key包括
  125. 'pid'所属项目id, 'tid'所属任务id,'name'已发布模型名称,
  126. 'path' 模型路径, 'exported_type' 已发布模型类型,
  127. """
  128. time_array = time.localtime(time.time())
  129. create_time = time.strftime("%Y-%m-%d %H:%M:%S", time_array)
  130. emid = workspace.max_exported_model_id + 1
  131. workspace.max_exported_model_id = emid
  132. if emid < 10000:
  133. emid = 'EM%04d' % emid
  134. else:
  135. emid = 'EM{}'.format(emid)
  136. pid = data['pid']
  137. tid = data['tid']
  138. name = data['name']
  139. path = data['path']
  140. exported_type = data['exported_type']
  141. assert pid in workspace.projects, "【已发布模型创建】项目ID'{}'不存在.".format(pid)
  142. assert tid in workspace.tasks, "【已发布模型创建】任务ID'{}'不存在.".format(tid)
  143. assert emid not in workspace.exported_models, "【已发布模型创建】已发布模型'{}'已经被占用.".format(
  144. emid)
  145. #assert osp.exists(path), "已发布模型路径不存在: {}".format(path)
  146. if not osp.exists(path):
  147. os.makedirs(path)
  148. task_path = workspace.tasks[tid].path
  149. # 拷贝评估结果
  150. eval_res_path = osp.join(task_path, 'eval_res.pkl')
  151. if osp.exists(eval_res_path):
  152. shutil.copy(eval_res_path, path)
  153. # 拷贝训练参数文件
  154. task_params_path = osp.join(task_path, 'params.pkl')
  155. if osp.exists(task_params_path):
  156. shutil.copy(task_params_path, path)
  157. # 拷贝数据集信息文件
  158. did = workspace.projects[pid].did
  159. dataset_path = workspace.datasets[did].path
  160. dataset_info_path = osp.join(dataset_path, "statis.pkl")
  161. if osp.exists(dataset_info_path):
  162. # 写入部分数据集信息
  163. with open(dataset_info_path, "rb") as f:
  164. dataset_info_dict = pickle.load(f)
  165. dataset_info_dict['name'] = workspace.datasets[did].name
  166. dataset_info_dict['desc'] = workspace.datasets[did].desc
  167. with open(dataset_info_path, "wb") as f:
  168. pickle.dump(dataset_info_dict, f)
  169. shutil.copy(dataset_info_path, path)
  170. from .project.task import get_task_params
  171. params = {'tid': tid}
  172. ret = get_task_params(params, workspace)
  173. train_params = ret['train']
  174. model_structure = train_params.model
  175. if hasattr(train_params, "backbone"):
  176. model_structure = "{}-{}".format(model_structure,
  177. train_params.backbone)
  178. if hasattr(train_params, "with_fpn"):
  179. if train_params.with_fpn:
  180. model_structure = "{}-{}".format(model_structure, "WITH_FPN")
  181. em = w.ExportedModel(
  182. id=emid,
  183. name=name,
  184. model=model_structure,
  185. type=workspace.projects[pid].type,
  186. pid=pid,
  187. tid=tid,
  188. create_time=create_time,
  189. path=path,
  190. exported_type=exported_type)
  191. workspace.exported_models[emid].CopyFrom(em)
  192. return {'status': 1, 'emid': emid}
  193. def list_exported_models(workspace):
  194. """列出预训练模型列表,可根据request中的参数进行筛选
  195. Args:
  196. """
  197. exported_model_list = list()
  198. for id in workspace.exported_models:
  199. exported_model = workspace.exported_models[id]
  200. model_id = exported_model.id
  201. model_name = exported_model.name
  202. model_model = exported_model.model
  203. model_type = exported_model.type
  204. model_pid = exported_model.pid
  205. model_tid = exported_model.tid
  206. model_create_time = exported_model.create_time
  207. model_path = exported_model.path
  208. model_exported_type = exported_model.exported_type
  209. attr = {
  210. 'id': model_id,
  211. 'name': model_name,
  212. 'model': model_model,
  213. 'type': model_type,
  214. 'pid': model_pid,
  215. 'tid': model_tid,
  216. 'create_time': model_create_time,
  217. 'path': model_path,
  218. 'exported_type': model_exported_type
  219. }
  220. if model_tid in workspace.tasks:
  221. from .project.task import get_export_status
  222. params = {'tid': model_tid}
  223. results = get_export_status(params, workspace)
  224. if results['export_status'] == TaskStatus.XEXPORTED:
  225. exported_model_list.append(attr)
  226. else:
  227. exported_model_list.append(attr)
  228. return {'status': 1, "exported_models": exported_model_list}
  229. def delete_exported_model(data, workspace):
  230. """删除exported_model。
  231. Args:
  232. data为dict,
  233. key包括'emid'已发布模型id
  234. """
  235. emid = data['emid']
  236. assert emid in workspace.exported_models, "已发布模型模型ID'{}'不存在.".format(emid)
  237. if osp.exists(workspace.exported_models[emid].path):
  238. shutil.rmtree(workspace.exported_models[emid].path)
  239. del workspace.exported_models[emid]
  240. return {'status': 1}
  241. def get_model_details(data, workspace):
  242. """获取模型详情。
  243. Args:
  244. data为dict,
  245. key包括'mid'模型id
  246. """
  247. mid = data['mid']
  248. if mid in workspace.pretrained_models:
  249. model_path = workspace.pretrained_models[mid].path
  250. elif mid in workspace.exported_models:
  251. model_path = workspace.exported_models[mid].path
  252. else:
  253. raise "模型{}不存在".format(mid)
  254. dataset_file = osp.join(model_path, 'statis.pkl')
  255. dataset_info = pickle.load(open(dataset_file, 'rb'))
  256. dataset_attr = {
  257. 'name': dataset_info['name'],
  258. 'desc': dataset_info['desc'],
  259. 'labels': dataset_info['labels'],
  260. 'train_num': len(dataset_info['train_files']),
  261. 'val_num': len(dataset_info['val_files']),
  262. 'test_num': len(dataset_info['test_files'])
  263. }
  264. task_params_file = osp.join(model_path, 'params.pkl')
  265. task_params = pickle.load(open(task_params_file, 'rb'))
  266. eval_result_file = osp.join(model_path, 'eval_res.pkl')
  267. eval_result = pickle.load(open(eval_result_file, 'rb'))
  268. #model_file = {'task_attr': task_params_file, 'eval_result': eval_result_file}
  269. return {
  270. 'status': 1,
  271. 'dataset_attr': dataset_attr,
  272. 'task_params': task_params,
  273. 'eval_result': eval_result
  274. }