[Pytorch] Airflow 사용 시 error: RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method 에러

 

상황

 

airflow의 GPU를 사용하는 DAG 코드를 작성하다가 다음과 같은 에러를 마주쳤다.

 

코드

from airflow                          import DAG
from airflow.operators.python         import PythonOperator
from datetime                         import datetime
from models                           import our_model
import torch

test = DAG(
    dag_id = 'test',
    start_date = datetime(2022,5,5),
    catchup = False,
    schedule = '@once'
)

train = PythonOperator(
    task_id         = 'train',
    python_callable = train,
    dag             = test
)

def train():
	device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model = our_model().to(device)

 

 

에러메시지

 

error: RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method

 

해결방법

 

device = 'cuda' if torch.cuda.is_available() else 'cpu'

 

이 부분을

 

device = 'cuda'

 

와 같이 수정해주면 된다.

 

이유는 모르겠으나 torch.cuda.is_available()을 하게 되면 저 에러 메시지가 계속 발생한다. GPU를 사용가능하다면 그냥 .to('cuda')를 사용해주면 된다.