parent
745d7516b3
commit
98204e2726
@ -1,88 +0,0 @@
|
|||||||
import torch
|
|
||||||
import os
|
|
||||||
from core import model_song # 학습할 때 썼던 model 파일을 불러와야 합니다.
|
|
||||||
|
|
||||||
# ----------------------------
|
|
||||||
# 1. 설정 (경로 및 입력 사이즈)
|
|
||||||
# ----------------------------
|
|
||||||
# 사용자님이 알려주신 ckpt 경로
|
|
||||||
ckpt_path = '/home/cuuva/face_exp/MobileFaceNet_Pytorch/result(MODEL_SONG)/MODEL_2_20251203_173900/best_model/best_001.ckpt'
|
|
||||||
onnx_path = 'model_2_test.onnx' # 저장될 파일 이름
|
|
||||||
|
|
||||||
# [중요] 학습할 때 사용한 이미지 해상도와 일치해야 합니다.
|
|
||||||
# 아까 코드에서 128x128로 수정하신 것을 확인했으므로 128로 설정합니다.
|
|
||||||
input_size = (1, 3, 128, 128)
|
|
||||||
|
|
||||||
def convert():
|
|
||||||
print(f"Loading checkpoint from: {ckpt_path}")
|
|
||||||
|
|
||||||
# ----------------------------
|
|
||||||
# 2. 모델 구조 정의
|
|
||||||
# ----------------------------
|
|
||||||
# 학습 코드와 동일한 모델 클래스를 인스턴스화 합니다.
|
|
||||||
net = model_song.MobileFacenet()
|
|
||||||
|
|
||||||
# ----------------------------
|
|
||||||
# 3. 가중치(Weights) 로드
|
|
||||||
# ----------------------------
|
|
||||||
checkpoint = torch.load(ckpt_path, map_location='cpu', weights_only=False) # GPU가 없어도 돌 수 있게 cpu로 로드
|
|
||||||
|
|
||||||
# 저장된 ckpt 구조에 따라 state_dict를 가져옵니다.
|
|
||||||
if 'net_state_dict' in checkpoint:
|
|
||||||
state_dict = checkpoint['net_state_dict']
|
|
||||||
else:
|
|
||||||
state_dict = checkpoint
|
|
||||||
|
|
||||||
# [핵심] DataParallel로 학습했다면 키(key) 앞에 'module.'이 붙어있습니다.
|
|
||||||
# 이를 제거해줘야 단일 모델에 로드할 수 있습니다.
|
|
||||||
new_state_dict = {}
|
|
||||||
for k, v in state_dict.items():
|
|
||||||
name = k.replace("module.", "") # 'module.conv1.weight' -> 'conv1.weight'
|
|
||||||
new_state_dict[name] = v
|
|
||||||
|
|
||||||
# 가중치 덮어씌우기
|
|
||||||
net.load_state_dict(new_state_dict)
|
|
||||||
|
|
||||||
# ----------------------------
|
|
||||||
# 4. 평가 모드 전환 (필수!)
|
|
||||||
# ----------------------------
|
|
||||||
# Dropout이나 Batch Norm이 학습 모드가 아닌 추론 모드로 동작하게 합니다.
|
|
||||||
net.eval()
|
|
||||||
# ----------------------------
|
|
||||||
# 4. ONNX 폴더 경로 생성
|
|
||||||
# ----------------------------
|
|
||||||
# ckpt_path 상위 폴더 이름 추출
|
|
||||||
experiment_folder_name = os.path.basename(os.path.dirname(os.path.dirname(ckpt_path)))
|
|
||||||
# 모델 최상위 경로
|
|
||||||
model_root = '/home/cuuva/face_exp/MobileFaceNet_Pytorch/result(MODEL_SONG)'
|
|
||||||
# 최종 ONNX 경로
|
|
||||||
onnx_dir = os.path.join(model_root, 'ONNX', experiment_folder_name)
|
|
||||||
os.makedirs(onnx_dir, exist_ok=True)
|
|
||||||
|
|
||||||
# ckpt 이름 기반으로 onnx 파일 이름 생성
|
|
||||||
# onnx_name = os.path.splitext(os.path.basename(ckpt_path))[0] + '.onnx'
|
|
||||||
onnx_name = 'model_2_test.onnx'
|
|
||||||
onnx_path = os.path.join(onnx_dir, onnx_name)
|
|
||||||
|
|
||||||
# ----------------------------
|
|
||||||
# 5. ONNX Export
|
|
||||||
# ----------------------------
|
|
||||||
print("Exporting to ONNX...")
|
|
||||||
|
|
||||||
# 모델 추적(Trace)을 위한 더미 입력 데이터 생성
|
|
||||||
dummy_input = torch.randn(*input_size)
|
|
||||||
|
|
||||||
torch.onnx.export(
|
|
||||||
net, # 실행할 모델
|
|
||||||
dummy_input, # 더미 입력값
|
|
||||||
onnx_path, # 저장할 경로
|
|
||||||
verbose=True, # 변환 과정 로그 출력
|
|
||||||
input_names=['input'], # 입력 노드 이름 (나중에 추론할 때 씀)
|
|
||||||
output_names=['output'], # 출력 노드 이름
|
|
||||||
external_data=False
|
|
||||||
)
|
|
||||||
|
|
||||||
print(f"Success! Model saved to: {os.path.abspath(onnx_path)}")
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
convert()
|
|
||||||
@ -1,166 +0,0 @@
|
|||||||
import os
|
|
||||||
import torch.utils.data
|
|
||||||
from torch import nn
|
|
||||||
from torch.nn import DataParallel
|
|
||||||
from datetime import datetime
|
|
||||||
from config import BATCH_SIZE, SAVE_FREQ, RESUME, SAVE_DIR, TEST_FREQ, TOTAL_EPOCH, MODEL_PRE, GPU
|
|
||||||
from config import CASIA_DATA_DIR, LFW_DATA_DIR
|
|
||||||
from core import model
|
|
||||||
from core.utils import init_log
|
|
||||||
from dataloader.CASIA_Face_loader import CASIA_Face
|
|
||||||
from dataloader.LFW_loader import LFW
|
|
||||||
from torch.optim import lr_scheduler
|
|
||||||
import torch.optim as optim
|
|
||||||
import time
|
|
||||||
from lfw_eval import parseList, evaluation_10_fold
|
|
||||||
import numpy as np
|
|
||||||
import scipy.io
|
|
||||||
|
|
||||||
# gpu init
|
|
||||||
gpu_list = ''
|
|
||||||
multi_gpus = False
|
|
||||||
if isinstance(GPU, int):
|
|
||||||
gpu_list = str(GPU)
|
|
||||||
else:
|
|
||||||
multi_gpus = True
|
|
||||||
for i, gpu_id in enumerate(GPU):
|
|
||||||
gpu_list += str(gpu_id)
|
|
||||||
if i != len(GPU) - 1:
|
|
||||||
gpu_list += ','
|
|
||||||
os.environ['CUDA_VISIBLE_DEVICES'] = gpu_list
|
|
||||||
|
|
||||||
# other init
|
|
||||||
start_epoch = 1
|
|
||||||
save_dir = os.path.join(SAVE_DIR, MODEL_PRE + 'v2_' + datetime.now().strftime('%Y%m%d_%H%M%S'))
|
|
||||||
if os.path.exists(save_dir):
|
|
||||||
raise NameError('model dir exists!')
|
|
||||||
os.makedirs(save_dir)
|
|
||||||
logging = init_log(save_dir)
|
|
||||||
_print = logging.info
|
|
||||||
|
|
||||||
|
|
||||||
# define trainloader and testloader
|
|
||||||
trainset = CASIA_Face(root=CASIA_DATA_DIR)
|
|
||||||
trainloader = torch.utils.data.DataLoader(trainset, batch_size=BATCH_SIZE,
|
|
||||||
shuffle=True, num_workers=8, drop_last=False)
|
|
||||||
|
|
||||||
# nl: left_image_path
|
|
||||||
# nr: right_image_path
|
|
||||||
nl, nr, folds, flags = parseList(root=LFW_DATA_DIR)
|
|
||||||
testdataset = LFW(nl, nr)
|
|
||||||
testloader = torch.utils.data.DataLoader(testdataset, batch_size=32,
|
|
||||||
shuffle=False, num_workers=8, drop_last=False)
|
|
||||||
|
|
||||||
# define model
|
|
||||||
net = model.MobileFacenet()
|
|
||||||
ArcMargin = model.ArcMarginProduct(128, trainset.class_nums)
|
|
||||||
|
|
||||||
if RESUME:
|
|
||||||
ckpt = torch.load(RESUME)
|
|
||||||
net.load_state_dict(ckpt['net_state_dict'])
|
|
||||||
start_epoch = ckpt['epoch'] + 1
|
|
||||||
|
|
||||||
|
|
||||||
# define optimizers
|
|
||||||
ignored_params = list(map(id, net.linear1.parameters()))
|
|
||||||
ignored_params += list(map(id, ArcMargin.weight))
|
|
||||||
prelu_params_id = []
|
|
||||||
prelu_params = []
|
|
||||||
for m in net.modules():
|
|
||||||
if isinstance(m, nn.PReLU):
|
|
||||||
ignored_params += list(map(id, m.parameters()))
|
|
||||||
prelu_params += m.parameters()
|
|
||||||
base_params = filter(lambda p: id(p) not in ignored_params, net.parameters())
|
|
||||||
|
|
||||||
optimizer_ft = optim.SGD([
|
|
||||||
{'params': base_params, 'weight_decay': 4e-5},
|
|
||||||
{'params': net.linear1.parameters(), 'weight_decay': 4e-4},
|
|
||||||
{'params': ArcMargin.weight, 'weight_decay': 4e-4},
|
|
||||||
{'params': prelu_params, 'weight_decay': 0.0}
|
|
||||||
], lr=0.1, momentum=0.9, nesterov=True)
|
|
||||||
|
|
||||||
exp_lr_scheduler = lr_scheduler.MultiStepLR(optimizer_ft, milestones=[36, 52, 58], gamma=0.1)
|
|
||||||
|
|
||||||
|
|
||||||
net = net.cuda()
|
|
||||||
ArcMargin = ArcMargin.cuda()
|
|
||||||
if multi_gpus:
|
|
||||||
net = DataParallel(net)
|
|
||||||
ArcMargin = DataParallel(ArcMargin)
|
|
||||||
criterion = torch.nn.CrossEntropyLoss()
|
|
||||||
|
|
||||||
|
|
||||||
best_acc = 0.0
|
|
||||||
best_epoch = 0
|
|
||||||
for epoch in range(start_epoch, TOTAL_EPOCH+1):
|
|
||||||
exp_lr_scheduler.step()
|
|
||||||
# train model
|
|
||||||
_print('Train Epoch: {}/{} ...'.format(epoch, TOTAL_EPOCH))
|
|
||||||
net.train()
|
|
||||||
|
|
||||||
train_total_loss = 0.0
|
|
||||||
total = 0
|
|
||||||
since = time.time()
|
|
||||||
for data in trainloader:
|
|
||||||
img, label = data[0].cuda(), data[1].cuda()
|
|
||||||
batch_size = img.size(0)
|
|
||||||
optimizer_ft.zero_grad()
|
|
||||||
|
|
||||||
raw_logits = net(img)
|
|
||||||
|
|
||||||
output = ArcMargin(raw_logits, label)
|
|
||||||
total_loss = criterion(output, label)
|
|
||||||
total_loss.backward()
|
|
||||||
optimizer_ft.step()
|
|
||||||
|
|
||||||
train_total_loss += total_loss.item() * batch_size
|
|
||||||
total += batch_size
|
|
||||||
|
|
||||||
train_total_loss = train_total_loss / total
|
|
||||||
time_elapsed = time.time() - since
|
|
||||||
loss_msg = ' total_loss: {:.4f} time: {:.0f}m {:.0f}s'\
|
|
||||||
.format(train_total_loss, time_elapsed // 60, time_elapsed % 60)
|
|
||||||
_print(loss_msg)
|
|
||||||
|
|
||||||
# test model on lfw
|
|
||||||
if epoch % TEST_FREQ == 0:
|
|
||||||
net.eval()
|
|
||||||
featureLs = None
|
|
||||||
featureRs = None
|
|
||||||
_print('Test Epoch: {} ...'.format(epoch))
|
|
||||||
for data in testloader:
|
|
||||||
for i in range(len(data)):
|
|
||||||
data[i] = data[i].cuda()
|
|
||||||
res = [net(d).data.cpu().numpy() for d in data]
|
|
||||||
featureL = np.concatenate((res[0], res[1]), 1)
|
|
||||||
featureR = np.concatenate((res[2], res[3]), 1)
|
|
||||||
if featureLs is None:
|
|
||||||
featureLs = featureL
|
|
||||||
else:
|
|
||||||
featureLs = np.concatenate((featureLs, featureL), 0)
|
|
||||||
if featureRs is None:
|
|
||||||
featureRs = featureR
|
|
||||||
else:
|
|
||||||
featureRs = np.concatenate((featureRs, featureR), 0)
|
|
||||||
|
|
||||||
result = {'fl': featureLs, 'fr': featureRs, 'fold': folds, 'flag': flags}
|
|
||||||
# save tmp_result
|
|
||||||
scipy.io.savemat('./result/tmp_result.mat', result)
|
|
||||||
accs = evaluation_10_fold('./result/tmp_result.mat')
|
|
||||||
_print(' ave: {:.4f}'.format(np.mean(accs) * 100))
|
|
||||||
|
|
||||||
# save model
|
|
||||||
if epoch % SAVE_FREQ == 0:
|
|
||||||
msg = 'Saving checkpoint: {}'.format(epoch)
|
|
||||||
_print(msg)
|
|
||||||
if multi_gpus:
|
|
||||||
net_state_dict = net.module.state_dict()
|
|
||||||
else:
|
|
||||||
net_state_dict = net.state_dict()
|
|
||||||
if not os.path.exists(save_dir):
|
|
||||||
os.mkdir(save_dir)
|
|
||||||
torch.save({
|
|
||||||
'epoch': epoch,
|
|
||||||
'net_state_dict': net_state_dict},
|
|
||||||
os.path.join(save_dir, '%03d.ckpt' % epoch))
|
|
||||||
print('finishing training')
|
|
||||||
@ -1,201 +0,0 @@
|
|||||||
import torch
|
|
||||||
import torch.optim as optim
|
|
||||||
from torch.optim import lr_scheduler
|
|
||||||
from torch.nn import DataParallel, CrossEntropyLoss
|
|
||||||
from dataloader.MyHF_loader import CASIA_HF, LFW_Pairs
|
|
||||||
from core import model_song
|
|
||||||
from core.utils import init_log
|
|
||||||
import os, time, numpy as np, scipy.io
|
|
||||||
from datetime import datetime
|
|
||||||
from config import BATCH_SIZE, SAVE_FREQ, RESUME, SAVE_DIR, TEST_FREQ, TOTAL_EPOCH, MODEL_PRE, GPU
|
|
||||||
from sklearn.metrics.pairwise import cosine_similarity # [추가] 정확도 계산용
|
|
||||||
|
|
||||||
# ----------------------------
|
|
||||||
# [추가] 간단한 LFW 정확도 계산 함수
|
|
||||||
# ----------------------------
|
|
||||||
def calculate_accuracy(featureLs, featureRs, flags, thresholds=np.arange(0, 1, 0.01)):
|
|
||||||
# 1. 특징 벡터 정규화 (Normalize)
|
|
||||||
featureLs = featureLs / np.linalg.norm(featureLs, axis=1, keepdims=True)
|
|
||||||
featureRs = featureRs / np.linalg.norm(featureRs, axis=1, keepdims=True)
|
|
||||||
|
|
||||||
# 2. 코사인 유사도 계산 (Dot Product)
|
|
||||||
scores = np.sum(featureLs * featureRs, axis=1)
|
|
||||||
|
|
||||||
# 3. 최적의 임계값(Threshold) 찾기 및 정확도 계산
|
|
||||||
best_acc = 0
|
|
||||||
for t in thresholds:
|
|
||||||
# 유사도가 t보다 크면 '같은 사람(1)', 작으면 '다른 사람(0)'
|
|
||||||
preds = (scores > t).astype(int)
|
|
||||||
acc = np.mean(preds == flags)
|
|
||||||
if acc > best_acc:
|
|
||||||
best_acc = acc
|
|
||||||
return best_acc
|
|
||||||
|
|
||||||
# ----------------------------
|
|
||||||
# GPU 및 초기 설정 (기존 동일)
|
|
||||||
# ----------------------------
|
|
||||||
gpu_list = ''
|
|
||||||
multi_gpus = False
|
|
||||||
if isinstance(GPU, int):
|
|
||||||
gpu_list = str(GPU)
|
|
||||||
else:
|
|
||||||
multi_gpus = True
|
|
||||||
gpu_list = ','.join(map(str, GPU))
|
|
||||||
os.environ['CUDA_VISIBLE_DEVICES'] = gpu_list
|
|
||||||
|
|
||||||
start_epoch = 1
|
|
||||||
save_dir = os.path.join('./result(MODEL_SONG)', 'MODEL_2_' + datetime.now().strftime('%Y%m%d_%H%M%S'))
|
|
||||||
os.makedirs(save_dir, exist_ok=True)
|
|
||||||
logging = init_log(save_dir)
|
|
||||||
_print = logging.info
|
|
||||||
|
|
||||||
# ----------------------------
|
|
||||||
# Dataloader (기존 동일)
|
|
||||||
# ----------------------------
|
|
||||||
trainset = CASIA_HF()
|
|
||||||
trainloader = torch.utils.data.DataLoader(trainset, batch_size=BATCH_SIZE,
|
|
||||||
shuffle=True, num_workers=8, drop_last=False)
|
|
||||||
|
|
||||||
testset = LFW_Pairs()
|
|
||||||
testloader = torch.utils.data.DataLoader(testset, batch_size=32,
|
|
||||||
shuffle=False, num_workers=8, drop_last=False)
|
|
||||||
|
|
||||||
# ----------------------------
|
|
||||||
# Model & Optimizer (기존 동일)
|
|
||||||
# ----------------------------
|
|
||||||
net = model_song.MobileFacenet()
|
|
||||||
ArcMargin = model_song.ArcMarginProduct(128, trainset.dataset.features['label'].num_classes)
|
|
||||||
|
|
||||||
if RESUME:
|
|
||||||
ckpt = torch.load(RESUME)
|
|
||||||
net.load_state_dict(ckpt['net_state_dict'])
|
|
||||||
start_epoch = ckpt['epoch'] + 1
|
|
||||||
|
|
||||||
net = net.cuda()
|
|
||||||
ArcMargin = ArcMargin.cuda()
|
|
||||||
if multi_gpus:
|
|
||||||
net = DataParallel(net)
|
|
||||||
ArcMargin = DataParallel(ArcMargin)
|
|
||||||
|
|
||||||
criterion = CrossEntropyLoss()
|
|
||||||
|
|
||||||
ignored_params = list(map(id, ArcMargin.weight))
|
|
||||||
# prelu_params = [p for m in net.modules() if isinstance(m, torch.nn.PReLU) for p in m.parameters()]
|
|
||||||
base_params = filter(lambda p: id(p) not in ignored_params, net.parameters())
|
|
||||||
|
|
||||||
# 기존 아키텍처에서 prelu 삭제했었으니까 아래 optim에서도 삭제 처리
|
|
||||||
optimizer_ft = optim.SGD([
|
|
||||||
{'params': base_params, 'weight_decay': 4e-5},
|
|
||||||
{'params': ArcMargin.weight, 'weight_decay': 4e-4}
|
|
||||||
], lr=0.1, momentum=0.9, nesterov=True)
|
|
||||||
|
|
||||||
# optimizer_ft = optim.SGD([
|
|
||||||
# {'params': base_params, 'weight_decay': 4e-5},
|
|
||||||
# {'params': net.linear1.parameters(), 'weight_decay': 4e-4},
|
|
||||||
# {'params': ArcMargin.weight, 'weight_decay': 4e-4},
|
|
||||||
# {'params': prelu_params, 'weight_decay': 0.0}
|
|
||||||
# ], lr=0.1, momentum=0.9, nesterov=True)
|
|
||||||
|
|
||||||
# 여기도 Config에서 Epoch 숫자 수정할때마다 milestone도 같이 수정해줘야함.
|
|
||||||
exp_lr_scheduler = lr_scheduler.MultiStepLR(optimizer_ft, milestones=[240, 310, 400], gamma=0.1)
|
|
||||||
|
|
||||||
# ----------------------------
|
|
||||||
# [추가] Best Accuracy 기록 변수
|
|
||||||
# ----------------------------
|
|
||||||
best_lfw_acc = 0.0
|
|
||||||
|
|
||||||
# ----------------------------
|
|
||||||
# Training Loop
|
|
||||||
# ----------------------------
|
|
||||||
for epoch in range(start_epoch, TOTAL_EPOCH + 1):
|
|
||||||
net.train()
|
|
||||||
train_total_loss, total = 0, 0
|
|
||||||
since = time.time()
|
|
||||||
_print(f"Train Epoch: {epoch}/{TOTAL_EPOCH} ...")
|
|
||||||
|
|
||||||
for data in trainloader:
|
|
||||||
img, label = data[0].cuda(), data[1].cuda()
|
|
||||||
optimizer_ft.zero_grad()
|
|
||||||
raw_logits = net(img)
|
|
||||||
output = ArcMargin(raw_logits, label)
|
|
||||||
loss = criterion(output, label)
|
|
||||||
loss.backward()
|
|
||||||
optimizer_ft.step()
|
|
||||||
train_total_loss += loss.item() * img.size(0)
|
|
||||||
total += img.size(0)
|
|
||||||
|
|
||||||
train_total_loss /= total
|
|
||||||
time_elapsed = time.time() - since
|
|
||||||
_print(f" total_loss: {train_total_loss:.4f} time: {time_elapsed//60:.0f}m {time_elapsed%60:.0f}s")
|
|
||||||
|
|
||||||
exp_lr_scheduler.step()
|
|
||||||
|
|
||||||
# ----------------------------
|
|
||||||
# Test & Best Model Save
|
|
||||||
# ----------------------------
|
|
||||||
if epoch % TEST_FREQ == 0:
|
|
||||||
net.eval()
|
|
||||||
featureLs, featureRs = None, None
|
|
||||||
flags = [] # [추가] 정답(Label)을 저장할 리스트
|
|
||||||
|
|
||||||
_print(" Testing LFW...")
|
|
||||||
with torch.no_grad(): # [추가] 테스트 땐 기울기 계산 끔 (메모리 절약)
|
|
||||||
for data in testloader:
|
|
||||||
# data 구조: [images_list, label(flag)]라고 가정
|
|
||||||
# LFW_Pairs의 경우 data[1]이 보통 정답(1:같은사람, 0:다른사람)
|
|
||||||
|
|
||||||
# 이미지 GPU 이동
|
|
||||||
imgs = [d.cuda() for d in data[0]]
|
|
||||||
|
|
||||||
# 정답 라벨 수집 (numpy로 변환)
|
|
||||||
flags.append(data[1].numpy())
|
|
||||||
|
|
||||||
# 특징 추출
|
|
||||||
res = [net(d).data.cpu().numpy() for d in imgs]
|
|
||||||
|
|
||||||
featureL = np.concatenate((res[0], res[1]), 1)
|
|
||||||
featureR = np.concatenate((res[2], res[3]), 1)
|
|
||||||
|
|
||||||
featureLs = featureL if featureLs is None else np.concatenate((featureLs, featureL), 0)
|
|
||||||
featureRs = featureR if featureRs is None else np.concatenate((featureRs, featureR), 0)
|
|
||||||
|
|
||||||
# [추가] 정답 리스트 합치기
|
|
||||||
flags = np.concatenate(flags, 0)
|
|
||||||
|
|
||||||
# [추가] 정확도 계산
|
|
||||||
# 만약 scipy.io.savemat은 필요하면 유지, 아니면 삭제해도 됨
|
|
||||||
# result = {'fl': featureLs, 'fr': featureRs}
|
|
||||||
# scipy.io.savemat('./result/tmp_result.mat', result)
|
|
||||||
|
|
||||||
# 직접 정확도 계산 (함수 호출)
|
|
||||||
current_acc = calculate_accuracy(featureLs, featureRs, flags)
|
|
||||||
_print(f" LFW Acc: {current_acc*100:.2f}% (Best: {best_lfw_acc*100:.2f}%)")
|
|
||||||
|
|
||||||
# [핵심] Best Model 저장 (Loss가 아닌 Acc 기준)
|
|
||||||
if current_acc > best_lfw_acc:
|
|
||||||
best_lfw_acc = current_acc
|
|
||||||
state_dict = net.module.state_dict() if multi_gpus else net.state_dict()
|
|
||||||
|
|
||||||
best_dir = os.path.join(save_dir, 'best_model')
|
|
||||||
os.makedirs(best_dir, exist_ok=True)
|
|
||||||
|
|
||||||
best_path = os.path.join(best_dir, f'best_{epoch:03d}.ckpt')
|
|
||||||
torch.save(
|
|
||||||
{
|
|
||||||
'epoch': epoch,
|
|
||||||
'net_state_dict': state_dict,
|
|
||||||
'acc': best_lfw_acc
|
|
||||||
},
|
|
||||||
best_path
|
|
||||||
)
|
|
||||||
_print(f" ==> Best Model Saved! (Acc: {best_lfw_acc*100:.2f}%, Epoch: {epoch}))")
|
|
||||||
|
|
||||||
# ----------------------------
|
|
||||||
# Regular Save (백업용)
|
|
||||||
# ----------------------------
|
|
||||||
if epoch % SAVE_FREQ == 0:
|
|
||||||
state_dict = net.module.state_dict() if multi_gpus else net.state_dict()
|
|
||||||
torch.save({'epoch': epoch, 'net_state_dict': state_dict},
|
|
||||||
os.path.join(save_dir, f'{epoch:03d}.ckpt'))
|
|
||||||
|
|
||||||
_print("finishing training")
|
|
||||||
Loading…
Reference in new issue