diff --git a/config/photo_review.py b/config/photo_review.py index 4dcbb17..c1d2318 100644 --- a/config/photo_review.py +++ b/config/photo_review.py @@ -15,3 +15,6 @@ IE_BATCH_SIZE = 4 # 是否开启布局分析 LAYOUT_ANALYSIS = False + +# 可使用的GPU +CUDA_VISIBLE_DEVICES = "1" diff --git a/main.py b/main.py index 660a77d..fa2bcc3 100644 --- a/main.py +++ b/main.py @@ -6,10 +6,10 @@ import paddle from auto_email.error_email import send_an_error_email from config.log import LOGGING_CONFIG -from config.photo_review import RETRY_TIME, SEND_ERROR_EMAIL +from config.photo_review import RETRY_TIME, SEND_ERROR_EMAIL, CUDA_VISIBLE_DEVICES from photo_review.photo_review import main -os.environ["CUDA_VISIBLE_DEVICES"] = "1" +os.environ["CUDA_VISIBLE_DEVICES"] = CUDA_VISIBLE_DEVICES # 项目必须从此处启动,否则代码中的相对路径可能导致错误的发生 if __name__ == '__main__': diff --git a/photo_review/__init__.py b/photo_review/__init__.py index 01fea81..44a965c 100644 --- a/photo_review/__init__.py +++ b/photo_review/__init__.py @@ -1,8 +1,12 @@ +import os + from paddlenlp import Taskflow from config.keys import PATIENT_NAME, ADMISSION_DATE, DISCHARGE_DATE, MEDICAL_EXPENSES, PERSONAL_CASH_PAYMENT, \ PERSONAL_ACCOUNT_PAYMENT, PERSONAL_FUNDED_AMOUNT, MEDICAL_INSURANCE_TYPE, HOSPITAL, DEPARTMENT, DOCTOR -from config.photo_review import IE_BATCH_SIZE, LAYOUT_ANALYSIS +from config.photo_review import IE_BATCH_SIZE, LAYOUT_ANALYSIS, CUDA_VISIBLE_DEVICES + +os.environ["CUDA_VISIBLE_DEVICES"] = CUDA_VISIBLE_DEVICES # 基本医保结算单 settlement_list_schema = \