반응형
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")
반응형
'취미 & 아무거나' 카테고리의 다른 글
[Compiler] Liveness analysis w. CFG (0) | 2024.06.17 |
---|---|
[Sensor] TF-LC02 LiDAR sensor Quickstart (0) | 2024.06.11 |
[YOLOv8] 이미지 확인을 위한 Image Convert (0) | 2024.06.04 |
[YOLOv8] YOLOv8n 맛보기 + 커스텀 데이터셋 (0) | 2024.06.03 |
[Keyboard] Nuphy TES68 (0) | 2024.05.13 |