AI/Naver_Boostcamp AI Tech

CPU와 GPU에서의 Tensor

john8538 2024. 8. 6. 16:07

목차

    CPU와 GPU 정리 & 차이점

    CPU란 중앙처리 장치를 뜻하며 컴퓨터 시스템을 통제하고 프로그램의 연산을 실행 - 처리하는 가장 핵심적인 컴퓨터 제어장치를 말한다. 수십억 개의 트랜지스터로 구축되었으며 여러 개의 프로세싱 코어를 갖출 수 있다.

     

    GPU란 더 작고 전문화된 코어로 구성된 프로세서다. 처리 작업을 병렬로 코어 간에 분할할 수 있으므로 대규모 성능을 제공한다.

     

    둘다 중요한 컴퓨팅 엔진이고 실리콘 기반 마이크로프로세서이며 데이터를 처리하나 중요한 점은 아키텍처가 크게 다르며 만들어진 용도가 다르다.

     

     

     

    주요 차이점은 아래 표에서 확인하자.

     

    Tensor에서의 CPU와 GPU

    이제 이 포스팅을 하게된 본론으로 들어가보자

    기본적으로 먼저 텐서를 생성한다.

    그 이후 텐서가 현재 어떤 디바이스에서 있는지 확인하려면 

    a = torch.tensor([1, 2, 3])
    a.device

     

    이와 같은 내용을 입력한다.

     

    이렇게 된다면 거의 99%의 경우네는 대부분 CPU 상태로 나올 것이다.

    이후 Pytorch 의 경우에 .to('cuda') 를 사용하여 GPU로 옮기곤 한다.

     

    이를 사용하여 옮기기전 확인해보아야 할점이 있다.

    torch.cuda_is_available()

     

    을 통하여 true가 출력 되는지 확인하여야 한다. 만약 나오지 않는경우 torch를 본인이 가지고 있는 gpu에 맞게 재설치를 진행해주어야 한다.

     

    애초에 탠서를 생성할때 gpu에 올리고 시작하고 싶다면 

    x = torch.tensor([1., 2.], device="cuda")

    와 같이 텐서를 생성해주어도 된다.

     

    따라서 총 텐서를 gpu로 생성하는 방법이 3가지가 존재한다.

     

    x = torch.tensor([1., 2.], device="cuda") 
    x = torch.tensor([1., 2.]).cuda() 
    x = torch.tensor([1., 2.]).to("cuda")

     

    물론 이렇게 명시해주는 편이 편하다.

    하지만 기본적으로 모든 텐서를 GPU에 생성하도록 설정하는 방법이 없을까? 

     

    궁금증이 생겨 찾아보다 한가지를 발견할 수 있었다.

    import torch
    
    # CUDA가 사용 가능한지 확인
    if torch.cuda.is_available():
        # 기본 텐서 타입을 CUDA 텐서로 설정
        torch.set_default_tensor_type(torch.cuda.FloatTensor)
    
    # 이제 생성되는 모든 텐서는 기본적으로 GPU에 위치하게 됩니다
    tensor1 = torch.tensor([1, 2, 3])  # GPU에 생성됨
    tensor2 = torch.randn(3, 3)        # GPU에 생성됨

    이러한 .set_default_tensor_type을 설정해주게 된다면 전역적으로 GPU에 적용이 되게 된다.

    간혹 cpu에서 이뤄지는 연산과 gpu에서 이뤄지는 연산을 명확히 구분지어야 할 경우가 존재하기도 하는데

    이 경우를 제외하고는 이렇게 할당해주게되면 코드가 더 간결해지고 GPU 사용 및 관리를 조금 더 수월하게 할 수 있게된다.

     

    argument를 이용한 CPU와 GPU

    파이썬 코드를 모듈화하여 분리를 진행해 py 파일로 실행을 하게된다면 코드내에 설정을 해주는것보다 argument로 코드를 실행시킬때 작성하는게 조금 더 편하다.

     

    코드는 아래와 같다.

    # python main.py --cpu -> CPU에서 동작
    # python main.py -> GPU에서 동작
    
    import argparse
    import torch
    
    # 명령줄 인자를 파싱하기 위한 ArgumentParser 객체 생성
    parser = argparse.ArgumentParser()
    
    # '--cpu' 인자 추가: 이 인자가 사용되면 CPU 모드로 실행
    parser.add_argument('--cpu', action='store_true', help='run in cpu')
    
    # 명령줄 인자 파싱
    args = parser.parse_args()
    
    # 실행 장치(device) 설정
    if args.cpu:
        # '--cpu' 인자가 주어졌다면 CPU를 사용
        device = torch.device('cpu')
    else:
        # '--cpu' 인자가 없으면 CUDA(GPU)를 사용
        device = torch.device('cuda')
    
    # 텐서 생성 및 지정된 장치로 이동
    x = torch.tensor([1., 2.]).to(device)

     

     

    특정 GPU만 사용하기

    GPU 서버를 쓰게 된다면 하나의 GPU를 할당 받을때도 있지만 여러개의 GPU도 할당 받기도 한다.

    만약 a100의 GPU를 4개 할당 받았다고 가정을 하자.

    기본적으로 코드를 실행하게 된다면 가장 먼저 0번째의 GPU에 할당이 자동으로 되게 된다.

    이를 코드를 통해 설정하게 된다면 각각의 GPU마다 다른 파이썬 파일을 할당할 수 있다.

    # GPU 1을 사용하기
    export CUDA_VISIBLE_DEVICES=1
    python main.py

     

    이와 같이 말이다.

     

    위에 코드와 병합지어 한 코드로 작성해보자

    import argparse
    import torch
    import os
    
    parser = argparse.ArgumentParser()
    parser.add_argument('--cpu', action='store_true', help='run in cpu')
    parser.add_argument('--gpu', type=int, default=0, help='GPU device index')
    args = parser.parse_args()
    
    if args.cpu:
        device = torch.device('cpu')
    else:
        # 사용 가능한 GPU 개수 확인
        num_gpus = torch.cuda.device_count()
        if args.gpu < num_gpus:
            # 지정된 GPU 사용
            device = torch.device(f'cuda:{args.gpu}')
            torch.cuda.set_device(args.gpu)
        else:
            print(f"GPU {args.gpu} is not available. Using CPU instead.")
            device = torch.device('cpu')
    
    # 현재 사용 중인 장치 출력
    print(f"Using device: {device}")
    
    # 텐서 생성 및 지정된 장치로 이동
    x = torch.tensor([1., 2.]).to(device)
    
    # 현재 텐서가 위치한 장치 확인
    print(f"Tensor x is on device: {x.device}")

     

    # CPU 사용
    python main.py --cpu

    # 기본 GPU (보통 0번) 사용
    python main.py

    # 특정 GPU (예: 1번) 사용
    python main.py --gpu 1

     

    이를 콘솔에 입력해주면 실행할 수 있게된다.

    만약 자신이 여러개의 GPU를 병렬로 사용하고 싶다면 각 GPU의 연산을 걸어주어야하는데 병렬화 코드를 작성해주어야 한다.

     

    # 여러 GPU에서 모델 병렬화 model = nn.DataParallel(model)

     

    혹은 from torch.nn.parallel import DistributedDataParallel as DDP 이를 사용하여야 한다.

     

    간혹 GPU 연산을 진행하다 나눠지는 연산이 있을 수 있어 각 연산이 어디 GPU에서 할당되는지 잘 알아보아야 한다.