PyTorch Basics

import torch
import torchvision.models as models

การบันทึกและโหลด Model Weights (วิธีที่นิยมที่สุด)

โมเดลใน PyTorch จะเก็บพารามิเตอร์ที่เรียนรู้มาแล้วไว้ในพจนานุกรมภายในที่เรียกว่า state_dict เราสามารถบันทึกเฉพาะค่าเหล่านี้เก็บไว้ในไฟล์นามสกุล .pth ได้ผ่านคำสั่ง torch.save

# โหลดโมเดลมาตรฐาน (VGG16) ที่ยังไม่ได้เทรนมาเป็นตัวอย่าง
model = models.vgg16(weights='IMAGENET1K_V1')

# บันทึกเฉพาะ Weights
torch.save(model.state_dict(), 'model_weights.pth')
print("Saved PyTorch Model State to model_weights.pth")

Saved PyTorch Model State to model_weights.pth

การโหลด Weights กลับมาใช้

ในการโหลด Weights แบบ state_dict กลับมา เราจำเป็นต้องสร้าง Instance ของโมเดลแบบเดียวกันขึ้นมาก่อน (เหมือนสร้างร่างเปล่าๆ ขึ้นมารอรับความจำ) จากนั้นค่อยโหลดพารามิเตอร์ใส่เข้าไปด้วยคำสั่ง load_state_dict()

# 1. สร้างโครงสร้างโมเดลเปล่าๆ (ร่างเปล่า)
# เราตั้ง weights=None เพื่อไม่โหลด weights จากเน็ต
model = models.vgg16(weights=None) 

# 2. โหลดความจำที่เราเซฟไว้ใส่เข้าไป
model.load_state_dict(torch.load('model_weights.pth'))

การบันทึกและโหลดทั้งโมเดล (Entire Model)

หากเราไม่อยากมานั่งเขียนสร้างโครงสร้าง (Instantiate model class) ใหม่ทุกครั้งเวลาจะโหลด เราสามารถสั่งเซฟ ทั้งโครงสร้าง (Structure) และน้ำหนัก (Weights) ลงไปในไฟล์เดียวกันได้เลย โดยส่งตัวแปรโมเดลเข้าไปตรงๆ (ไม่ได้ส่งแค่ state_dict)

# บันทึกทั้งโมเดล
torch.save(model, 'model.pth')
print("Saved Full PyTorch Model to model.pth")

Saved Full PyTorch Model to model.pth

การโหลด Entire Model

การโหลดกลับมาก็ทำได้ง่ายๆ บรรทัดเดียว โดยไม่ต้องสร้างร่างเปล่ามารองรับ

# โหลดทั้งโมเดลกลับมา
model = torch.load('model.pth')

หมายเหตุ: การเซฟแบบ Entire Model นี้ ภายใต้เบื้องหลัง PyTorch ใช้โมดูล pickle ของ Python ดังนั้น การจะโหลด model.pth ขึ้นมาใช้งานได้ โค้ดโปรเจกต์ของคุณจะต้องสามารถเข้าถึงคลาสหรือโครงสร้างต้นฉบับได้ในขณะที่รัน (Class Definition must be available)