[YOLOv8] YOLOv8n 내부 가중치 확인하기

GPT 의 힘을 빌려서 완성했다.

command 창 내 현재 위치에 파일이 생성되니 주의할 것.

import torch
from ultralytics import YOLO

# ------------ 모델 경로 및 파일 수정 시간은 수정하기 --------------
model_path = '../best.pt' 
model_generated_time = "  24-06-04, 20:02"
# --------------------------------------------------------------

model = YOLO(model_path)

# 모델의 가중치 확인
model_weights = model.model.state_dict()
torch.set_printoptions(threshold=torch.inf)
idx = 0
while idx < len(model_weights):
    name = list(model_weights.keys())[idx]
    param = model_weights[name]
    if "model." in name:
        output_file = f"{name}.txt" 
        with open(output_file, "w") as f:
            f.write("Weights Information:\n")
            f.write("# weights ver.2\n\n")
            f.write(f"Model Path: {model_path}\n")
            f.write(f"Model Generated Time: {model_generated_time}\n\n")
            f.write(f"Layer: {name} | Size: {param.size()}\n")
            f.write(str(param) + "\n")  
    idx += 1

 

 

다음은 레이어 정보만 가져오는 코드임.

 

# 레이어 정보만 가져오기
import torch
from ultralytics import YOLO

# 모델 경로 및 파일 수정 시간 설정
model_path = '../best.pt' 
model_generated_time = "24-06-04, 20:02"

# YOLO 모델 로드
model = YOLO(model_path)

# 모델의 가중치 확인
model_weights = model.model.state_dict()

# 레이어 정보를 저장할 파일 생성
with open("layers.txt", "w") as f:
    f.write("Layer Information:\n")
    f.write(f"Model Path: {model_path}\n")
    f.write(f"Model Generated Time: {model_generated_time}\n\n")

# 각 레이어 정보를 파일에 기록
with open("layers.txt", "a") as f:
    for idx, (name, param) in enumerate(model_weights.items()):
        f.write(f"Layer {idx}:\n")
        f.write(f"Name: {name}\n")
        f.write(f"Size: {param.size()}\n")