提交 73a3415c 编写于 作者: H Hongyu Liu 提交者: pkpk

Add cuda check in language model (#2722)

* add gpu check; test=develop

* add cuda check in train; test=develop
上级 7e00cc08
...@@ -35,12 +35,12 @@ import sys ...@@ -35,12 +35,12 @@ import sys
if sys.version[0] == '2': if sys.version[0] == '2':
reload(sys) reload(sys)
sys.setdefaultencoding("utf-8") sys.setdefaultencoding("utf-8")
sys.path.append('..') sys.path.append('../')
import os import os
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
from args import * from args import *
sys.path.append("../") from models.model_check import check_cuda
from models.language_model import lm_model from models.language_model import lm_model
from config import RNNConfig from config import RNNConfig
import logging import logging
...@@ -87,6 +87,9 @@ def save_para_npz(train_prog, train_exe): ...@@ -87,6 +87,9 @@ def save_para_npz(train_prog, train_exe):
def main(): def main():
args = parse_args() args = parse_args()
check_cuda(args.use_gpu)
logger = logging.getLogger("lm") logger = logging.getLogger("lm")
logger.setLevel(logging.INFO) logger.setLevel(logging.INFO)
formatter = logging.Formatter( formatter = logging.Formatter(
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册