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