Yolov8 GPU로 학습시키기

 개요

참고한 사이트: https://artiiicy.tistory.com/61

[Pytorch] 특정 GPU 사용하기 / 여러개의 multi GPU parallel 하게 사용하기

 

참고한 사이트2: https://docs.ultralytics.com/ko/modes/train/#apple-m1-and-m2-mps-training

 

코드 1: GPU로 돌리기

코드 2: 멀티 GPU 데이터 병렬 모드 

 

 코드 1

import os
import torch
import torch.multiprocessing as mp
from ultralytics import YOLO

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"  # Arrange GPU devices starting from 0
os.environ["CUDA_VISIBLE_DEVICES"]= "2"  # Set the GPU 2 to use

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

print('Device:', device)
print('Current cuda device:', torch.cuda.current_device())
print('Count of using GPUs:', torch.cuda.device_count())

if __name__ == '__main__':
    #mp.set_start_method("spawn")

    # Load a model
    model = YOLO('yolov8x.pt')  # load a pretrained model (recommended for training)

    # Train
    results = model.train(data='cfg/datasets/SOD_Drone.yaml', epochs=10, imgsz=640, device='cuda')

 

코드2

import os
import torch
import torch.multiprocessing as mp
from ultralytics import YOLO

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"  # Arrange GPU devices starting from 0
os.environ["CUDA_VISIBLE_DEVICES"]= "3, 4, 5, 6"  # 3번, 4번, 5번, 6번 사용 

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

print('Device:', device)
print('Current cuda device:', torch.cuda.current_device())
print('Count of using GPUs:', torch.cuda.device_count())

if __name__ == '__main__':
    #mp.set_start_method("spawn")

    # Load a model
    model = YOLO('yolov8x.pt')  # load a pretrained model (recommended for training)

    # Train, device를 사용하고 싶은 GPU 이름을 넣어줘야 함 
    results = model.train(data='cfg/datasets/SOD_Drone.yaml', epochs=100, imgsz=640, device=[3,4,5,6])

 

 번외 (나만 볼듯?)

도커에서 돌리는 법 

# hsh_cont2라는 도커 터미널 실행시키는 방법
docker exec -it hsh_cont2 /bin/bash

#주피터 노트북 서버 백그라운드에서 돌리기 
nohup jupyter lab > log.txt 2>&1 &


# 도커를 통해서 GPU 백그라운드에서 돌리는 방법 log.log는 훈련중인 log 볼 수 있음
nohup python3 yolov8_train.py > log.log 2>&1 &

# 터미널에서 로그 확인하는 방법 
cat log.log
tail log.log
tail -n 100 log.log

# 리눅스 터미널에서 6번 GPU를 통해 python 코드 돌리는 방법 
CUDA_VISIBLE_DEVICES=6 python3 yolov8_train.py