TTA
/
Face_recognition
18
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Inital project upload

Browse Source
main
hgkim 7 months ago
parent 0f552c3e99
commit dda56b30e5
20 changed files with 2271 additions and 152 deletions
Show Stats Download Patch File Download Diff File

248
.gitignore vendored
Unescape View File

@ -1,162 +1,106 @@
# ---> Python
# Byte-compiled / optimized / DLL files
# -----------------------------
# OS 기본 파일
# -----------------------------
.DS_Store
Thumbs.db
# -----------------------------
# Python 환경
# -----------------------------
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.pyc
*.pyo
*.pyd
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
share/python-wheels/
# 가상환경
env/
venv/
.mipenv/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
cover/
.eggs/
# Translations
*.mo
*.pot
# Poetry / pipenv
.cache/
.venv/
# -----------------------------
# IDE / Editor 관련
# -----------------------------
.vscode/
.idea/
*.swp
*.swo
# -----------------------------
# 데이터, 모델, 체크포인트
# -----------------------------
# 모델 파일
*.pt
*.pth
*.ckpt
*.onnx
*.trt
*.pb
*.h5
# 학습 관련 출력
runs/
logs/
tensorboard/
lightning_logs/
checkpoint/
checkpoints/
# 데이터셋 (원하면 제외 가능)
data/
dataset/
datasets/
# 결과물 (이미지/비디오/추론 결과)
outputs/
results/
inference/
*.jpg
*.png
*.jpeg
*.bmp
*.mp4
*.avi
# -----------------------------
# PyTorch & HuggingFace 캐시
# -----------------------------
/root/.cache/torch/
/cache/
/.torch/
huggingface/
transformers/
# -----------------------------
# Jupyter 관련
# -----------------------------
.ipynb_checkpoints/
*.ipynb~
# -----------------------------
# 컴파일/빌드 아티팩트
# -----------------------------
build/
dist/
*.bin
# Django stuff:
# -----------------------------
# Temp 파일
# -----------------------------
tmp/
temp/
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
.pybuilder/
target/
# Jupyter Notebook
.ipynb_checkpoints
# IPython
profile_default/
ipython_config.py
# pyenv
# For a library or package, you might want to ignore these files since the code is
# intended to run in multiple environments; otherwise, check them in:
# .python-version
# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock
# poetry
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
# This is especially recommended for binary packages to ensure reproducibility, and is more
# commonly ignored for libraries.
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
#poetry.lock
# pdm
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
#pdm.lock
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
# in version control.
# https://pdm.fming.dev/#use-with-ide
.pdm.toml
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# SageMath parsed files
*.sage.py
# Environments
# -----------------------------
# 환경 변수 파일
# -----------------------------
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
# pytype static type analyzer
.pytype/
# Cython debug symbols
cython_debug/
# PyCharm
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
# and can be added to the global gitignore or merged into this file. For a more nuclear
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
#.idea/
.env.*

15
config.py
Unescape View File

@ -0,0 +1,15 @@
BATCH_SIZE = 256
SAVE_FREQ = 1
TEST_FREQ = 1
TOTAL_EPOCH = 500
RESUME = ''
SAVE_DIR = './model'
MODEL_PRE = 'CASIA_B512_'
CASIA_DATA_DIR = '/home/xiaocc/Documents/caffe_project/sphereface/train/data'
LFW_DATA_DIR = '/home/xiaocc/Documents/caffe_project/sphereface/test/data'
GPU = 0

195
core/model.py
Unescape View File

@ -0,0 +1,195 @@
from torch import nn
import torch
import torch.nn.functional as F
from torch.autograd import Variable
import math
from torch.nn import Parameter
class Bottleneck(nn.Module):
def __init__(self, inp, oup, stride, expansion):
super(Bottleneck, self).__init__()
self.connect = stride == 1 and inp == oup
self.conv = nn.Sequential(
#pw
nn.Conv2d(inp, inp * expansion, 1, 1, 0, bias=False),
nn.BatchNorm2d(inp * expansion),
nn.ReLU(inplace=True),
#dw
nn.Conv2d(inp * expansion, inp * expansion, 3, stride, 1, groups=inp * expansion, bias=False),
nn.BatchNorm2d(inp * expansion),
nn.ReLU(inplace=True),
#pw-linear
nn.Conv2d(inp * expansion, oup, 1, 1, 0, bias=False),
nn.BatchNorm2d(oup),
)
def forward(self, x):
if self.connect:
return x + self.conv(x)
else:
return self.conv(x)
# class ConvBlock(nn.Module): # prelu 버전
# def __init__(self, inp, oup, k, s, p, dw=False, linear=False):
# super(ConvBlock, self).__init__()
# self.linear = linear
# if dw:
# self.conv = nn.Conv2d(inp, oup, k, s, p, groups=inp, bias=False)
# else:
# self.conv = nn.Conv2d(inp, oup, k, s, p, bias=False)
# self.bn = nn.BatchNorm2d(oup)
# if not linear:
# self.prelu = nn.PReLU(oup)
# def forward(self, x):
# x = self.conv(x)
# x = self.bn(x)
# if self.linear:
# return x
# else:
# return self.prelu(x)
class ConvBlock(nn.Module):
def __init__(self, inp, oup, k, s, p, dw=False, linear=False):
super(ConvBlock, self).__init__()
self.linear = linear
if dw:
self.conv = nn.Conv2d(inp, oup, k, s, p, groups=inp, bias=False)
else:
self.conv = nn.Conv2d(inp, oup, k, s, p, bias=False)
self.bn = nn.BatchNorm2d(oup)
if not linear:
self.relu = nn.ReLU(inplace=True)
def forward(self, x):
x = self.conv(x)
x = self.bn(x)
if self.linear:
return x
else:
return self.relu(x)
class ConvBlockAvgPool(nn.Module): # 이게...맞나?
def __init__(self, kernel):
super().__init__()
self.pool = nn.AvgPool2d(kernel)
self.bn = nn.BatchNorm2d(512)
def forward(self, x):
x = self.pool(x)
return self.bn(x)
# return self.pool(x)
Mobilefacenet_bottleneck_setting = [
# t, c , n ,s
[2, 64, 5, 2],
[4, 128, 1, 2],
[2, 128, 6, 1],
[4, 128, 1, 2],
[2, 128, 2, 1]
]
Mobilenetv2_bottleneck_setting = [
# t, c, n, s
[1, 16, 1, 1],
[6, 24, 2, 2],
[6, 32, 3, 2],
[6, 64, 4, 2],
[6, 96, 3, 1],
[6, 160, 3, 2],
[6, 320, 1, 1],
]
class MobileFacenet(nn.Module):
def __init__(self, bottleneck_setting=Mobilefacenet_bottleneck_setting):
super(MobileFacenet, self).__init__()
self.conv1 = ConvBlock(3, 64, 3, 2, 1)
self.dw_conv1 = ConvBlock(64, 64, 3, 1, 1, dw=True)
self.inplanes = 64
block = Bottleneck
self.blocks = self._make_layer(block, bottleneck_setting)
self.conv2 = ConvBlock(128, 512, 1, 1, 0)
# self.linear7 = ConvBlock(512, 512, (7, 6), 1, 0, dw=True, linear=True)
# self.linear7 = ConvBlock(512, 512, (8, 8), 1, 0, dw=True, linear=True) # 128x128 로 키우니까 커널사이즈도 키워줘야함.
# self.linear7 = nn.AvgPool2d(kernel_size=8, stride=1) # 여기봐바 여기 너가 말한대로 추가해놨어.
self.linear7 = ConvBlockAvgPool(kernel=8)
self.linear1 = ConvBlock(512, 128, 1, 1, 0, linear=True)
for m in self.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
def _make_layer(self, block, setting):
layers = []
for t, c, n, s in setting:
for i in range(n):
if i == 0:
layers.append(block(self.inplanes, c, s, t))
else:
layers.append(block(self.inplanes, c, 1, t))
self.inplanes = c
return nn.Sequential(*layers)
def forward(self, x):
x = self.conv1(x)
x = self.dw_conv1(x)
x = self.blocks(x)
x = self.conv2(x)
x = self.linear7(x)
x = self.linear1(x) # 이때 shape이 [Batch,128,1,1] 임.
x = x.view(x.size(0), -1) # reshpape에 해당되는 부분
return x
class ArcMarginProduct(nn.Module):
def __init__(self, in_features=128, out_features=200, s=32.0, m=0.50, easy_margin=False):
super(ArcMarginProduct, self).__init__()
self.in_features = in_features
self.out_features = out_features
self.s = s
self.m = m
self.weight = Parameter(torch.Tensor(out_features, in_features))
nn.init.xavier_uniform_(self.weight)
# init.kaiming_uniform_()
# self.weight.data.normal_(std=0.001)
self.easy_margin = easy_margin
self.cos_m = math.cos(m)
self.sin_m = math.sin(m)
# make the function cos(theta+m) monotonic decreasing while theta in [0°,180°]
self.th = math.cos(math.pi - m)
self.mm = math.sin(math.pi - m) * m
def forward(self, x, label):
cosine = F.linear(F.normalize(x), F.normalize(self.weight))
sine = torch.sqrt(1.0 - torch.pow(cosine, 2))
phi = cosine * self.cos_m - sine * self.sin_m
if self.easy_margin:
phi = torch.where(cosine > 0, phi, cosine)
else:
phi = torch.where((cosine - self.th) > 0, phi, cosine - self.mm)
one_hot = torch.zeros(cosine.size(), device='cuda')
one_hot.scatter_(1, label.view(-1, 1).long(), 1)
output = (one_hot * phi) + ((1.0 - one_hot) * cosine)
output *= self.s
return output
if __name__ == "__main__":
# input = Variable(torch.FloatTensor(2, 3, 112, 96))
input = Variable(torch.FloatTensor(2, 3, 128, 128)) # 해상도 128x128 수정 진행.
net = MobileFacenet()
print(net)
x = net(input)
print(x.shape)

187
core/model2.py
Unescape View File

@ -0,0 +1,187 @@
from torch import nn
import torch
import torch.nn.functional as F
from torch.autograd import Variable
import math
from torch.nn import Parameter
class Bottleneck(nn.Module):
def __init__(self, inp, oup, stride, expansion):
super(Bottleneck, self).__init__()
self.connect = stride == 1 and inp == oup
self.conv = nn.Sequential(
#pw
nn.Conv2d(inp, inp * expansion, 1, 1, 0, bias=False),
nn.BatchNorm2d(inp * expansion),
nn.ReLU(inplace=True),
#dw
nn.Conv2d(inp * expansion, inp * expansion, 3, stride, 1, groups=inp * expansion, bias=False),
nn.BatchNorm2d(inp * expansion),
nn.ReLU(inplace=True),
#pw-linear
nn.Conv2d(inp * expansion, oup, 1, 1, 0, bias=False),
nn.BatchNorm2d(oup),
)
def forward(self, x):
if self.connect:
return x + self.conv(x)
else:
return self.conv(x)
# class ConvBlock(nn.Module): # prelu 버전
# def __init__(self, inp, oup, k, s, p, dw=False, linear=False):
# super(ConvBlock, self).__init__()
# self.linear = linear
# if dw:
# self.conv = nn.Conv2d(inp, oup, k, s, p, groups=inp, bias=False)
# else:
# self.conv = nn.Conv2d(inp, oup, k, s, p, bias=False)
# self.bn = nn.BatchNorm2d(oup)
# if not linear:
# self.prelu = nn.PReLU(oup)
# def forward(self, x):
# x = self.conv(x)
# x = self.bn(x)
# if self.linear:
# return x
# else:
# return self.prelu(x)
class ConvBlock(nn.Module):
def __init__(self, inp, oup, k, s, p, dw=False, linear=False):
super(ConvBlock, self).__init__()
self.linear = linear
if dw:
self.conv = nn.Conv2d(inp, oup, k, s, p, groups=inp, bias=False)
else:
self.conv = nn.Conv2d(inp, oup, k, s, p, bias=False)
self.bn = nn.BatchNorm2d(oup)
if not linear:
self.relu = nn.ReLU(inplace=True)
def forward(self, x):
x = self.conv(x)
x = self.bn(x)
if self.linear:
return x
else:
return self.relu(x)
Mobilefacenet_bottleneck_setting = [
# t, c , n ,s
[2, 64, 5, 2],
[4, 128, 1, 2],
[2, 128, 6, 1],
[4, 128, 1, 2],
[2, 128, 2, 1]
]
Mobilenetv2_bottleneck_setting = [
# t, c, n, s
[1, 16, 1, 1],
[6, 24, 2, 2],
[6, 32, 3, 2],
[6, 64, 4, 2],
[6, 96, 3, 1],
[6, 160, 3, 2],
[6, 320, 1, 1],
]
class MobileFacenet(nn.Module):
def __init__(self, bottleneck_setting=Mobilefacenet_bottleneck_setting):
super(MobileFacenet, self).__init__()
self.conv1 = ConvBlock(3, 64, 3, 2, 1)
self.dw_conv1 = ConvBlock(64, 64, 3, 1, 1, dw=True)
self.inplanes = 64
block = Bottleneck
self.blocks = self._make_layer(block, bottleneck_setting)
self.conv2 = ConvBlock(128, 512, 1, 1, 0)
# self.linear7 = ConvBlock(512, 512, (7, 6), 1, 0, dw=True, linear=True)
# self.linear7 = ConvBlock(512, 512, 8, 1, 0, dw=True, linear=True) # 128x128 로 키우니까 커널사이즈도 키워줘야함.
self.pool = nn.AdaptiveAvgPool2d(1)
self.pw_conv = nn.Conv2d(512, 512, 1, 1, 0, bias=False)
self.bn7 = nn.BatchNorm2d(512)
self.linear1 = ConvBlock(512, 128, 1, 1, 0, linear=True)
for m in self.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
def _make_layer(self, block, setting):
layers = []
for t, c, n, s in setting:
for i in range(n):
if i == 0:
layers.append(block(self.inplanes, c, s, t))
else:
layers.append(block(self.inplanes, c, 1, t))
self.inplanes = c
return nn.Sequential(*layers)
def forward(self, x):
x = self.conv1(x)
x = self.dw_conv1(x)
x = self.blocks(x)
x = self.conv2(x)
# x = self.linear7(x)
x = self.pool(x)
x = self.pw_conv(x)
x = self.bn7(x)
x = self.linear1(x)
x = x.view(x.size(0), -1)
return x
class ArcMarginProduct(nn.Module):
def __init__(self, in_features=128, out_features=200, s=32.0, m=0.50, easy_margin=False):
super(ArcMarginProduct, self).__init__()
self.in_features = in_features
self.out_features = out_features
self.s = s
self.m = m
self.weight = Parameter(torch.Tensor(out_features, in_features))
nn.init.xavier_uniform_(self.weight)
# init.kaiming_uniform_()
# self.weight.data.normal_(std=0.001)
self.easy_margin = easy_margin
self.cos_m = math.cos(m)
self.sin_m = math.sin(m)
# make the function cos(theta+m) monotonic decreasing while theta in [0°,180°]
self.th = math.cos(math.pi - m)
self.mm = math.sin(math.pi - m) * m
def forward(self, x, label):
cosine = F.linear(F.normalize(x), F.normalize(self.weight))
sine = torch.sqrt(1.0 - torch.pow(cosine, 2))
phi = cosine * self.cos_m - sine * self.sin_m
if self.easy_margin:
phi = torch.where(cosine > 0, phi, cosine)
else:
phi = torch.where((cosine - self.th) > 0, phi, cosine - self.mm)
one_hot = torch.zeros(cosine.size(), device='cuda')
one_hot.scatter_(1, label.view(-1, 1).long(), 1)
output = (one_hot * phi) + ((1.0 - one_hot) * cosine)
output *= self.s
return output
if __name__ == "__main__":
# input = Variable(torch.FloatTensor(2, 3, 112, 96))
input = Variable(torch.FloatTensor(2, 3, 128, 128)) # 해상도 128x128 수정 진행.
net = MobileFacenet()
print(net)
x = net(input)
print(x.shape)

182
core/model_bak.py
Unescape View File

@ -0,0 +1,182 @@
from torch import nn
import torch
import torch.nn.functional as F
from torch.autograd import Variable
import math
from torch.nn import Parameter
class Bottleneck(nn.Module):
def __init__(self, inp, oup, stride, expansion):
super(Bottleneck, self).__init__()
self.connect = stride == 1 and inp == oup
self.conv = nn.Sequential(
#pw
nn.Conv2d(inp, inp * expansion, 1, 1, 0, bias=False),
nn.BatchNorm2d(inp * expansion),
nn.ReLU(inplace=True),
#dw
nn.Conv2d(inp * expansion, inp * expansion, 3, stride, 1, groups=inp * expansion, bias=False),
nn.BatchNorm2d(inp * expansion),
nn.ReLU(inplace=True),
#pw-linear
nn.Conv2d(inp * expansion, oup, 1, 1, 0, bias=False),
nn.BatchNorm2d(oup),
)
def forward(self, x):
if self.connect:
return x + self.conv(x)
else:
return self.conv(x)
# class ConvBlock(nn.Module): # prelu 버전
# def __init__(self, inp, oup, k, s, p, dw=False, linear=False):
# super(ConvBlock, self).__init__()
# self.linear = linear
# if dw:
# self.conv = nn.Conv2d(inp, oup, k, s, p, groups=inp, bias=False)
# else:
# self.conv = nn.Conv2d(inp, oup, k, s, p, bias=False)
# self.bn = nn.BatchNorm2d(oup)
# if not linear:
# self.prelu = nn.PReLU(oup)
# def forward(self, x):
# x = self.conv(x)
# x = self.bn(x)
# if self.linear:
# return x
# else:
# return self.prelu(x)
class ConvBlock(nn.Module):
def __init__(self, inp, oup, k, s, p, dw=False, linear=False):
super(ConvBlock, self).__init__()
self.linear = linear
if dw:
self.conv = nn.Conv2d(inp, oup, k, s, p, groups=inp, bias=False)
else:
self.conv = nn.Conv2d(inp, oup, k, s, p, bias=False)
self.bn = nn.BatchNorm2d(oup)
if not linear:
self.relu = nn.ReLU(inplace=True)
def forward(self, x):
x = self.conv(x)
x = self.bn(x)
if self.linear:
return x
else:
return self.relu(x)
Mobilefacenet_bottleneck_setting = [
# t, c , n ,s
[2, 64, 5, 2],
[4, 128, 1, 2],
[2, 128, 6, 1],
[4, 128, 1, 2],
[2, 128, 2, 1]
]
Mobilenetv2_bottleneck_setting = [
# t, c, n, s
[1, 16, 1, 1],
[6, 24, 2, 2],
[6, 32, 3, 2],
[6, 64, 4, 2],
[6, 96, 3, 1],
[6, 160, 3, 2],
[6, 320, 1, 1],
]
class MobileFacenet(nn.Module):
def __init__(self, bottleneck_setting=Mobilefacenet_bottleneck_setting):
super(MobileFacenet, self).__init__()
self.conv1 = ConvBlock(3, 64, 3, 2, 1)
self.dw_conv1 = ConvBlock(64, 64, 3, 1, 1, dw=True)
self.inplanes = 64
block = Bottleneck
self.blocks = self._make_layer(block, bottleneck_setting)
self.conv2 = ConvBlock(128, 512, 1, 1, 0)
# self.linear7 = ConvBlock(512, 512, (7, 6), 1, 0, dw=True, linear=True)
self.linear7 = ConvBlock(512, 512, 8, 1, 0, dw=True, linear=True) # (8,8) 안하고 8 하니까 이것도 loss 안주는듯? 아니다아니다
# self.linear7 = ConvBlock(512, 512, (8,8), 1, 0, dw=True, linear=True) # 128x128 로 키우니까 커널사이즈도 키워줘야함.
self.linear1 = ConvBlock(512, 128, 1, 1, 0, linear=True)
for m in self.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
def _make_layer(self, block, setting):
layers = []
for t, c, n, s in setting:
for i in range(n):
if i == 0:
layers.append(block(self.inplanes, c, s, t))
else:
layers.append(block(self.inplanes, c, 1, t))
self.inplanes = c
return nn.Sequential(*layers)
def forward(self, x):
x = self.conv1(x)
x = self.dw_conv1(x)
x = self.blocks(x)
x = self.conv2(x)
x = self.linear7(x)
x = self.linear1(x)
x = x.view(x.size(0), -1)
return x
class ArcMarginProduct(nn.Module):
def __init__(self, in_features=128, out_features=200, s=32.0, m=0.50, easy_margin=False):
super(ArcMarginProduct, self).__init__()
self.in_features = in_features
self.out_features = out_features
self.s = s
self.m = m
self.weight = Parameter(torch.Tensor(out_features, in_features))
nn.init.xavier_uniform_(self.weight)
# init.kaiming_uniform_()
# self.weight.data.normal_(std=0.001)
self.easy_margin = easy_margin
self.cos_m = math.cos(m)
self.sin_m = math.sin(m)
# make the function cos(theta+m) monotonic decreasing while theta in [0°,180°]
self.th = math.cos(math.pi - m)
self.mm = math.sin(math.pi - m) * m
def forward(self, x, label):
cosine = F.linear(F.normalize(x), F.normalize(self.weight))
sine = torch.sqrt(1.0 - torch.pow(cosine, 2))
phi = cosine * self.cos_m - sine * self.sin_m
if self.easy_margin:
phi = torch.where(cosine > 0, phi, cosine)
else:
phi = torch.where((cosine - self.th) > 0, phi, cosine - self.mm)
one_hot = torch.zeros(cosine.size(), device='cuda')
one_hot.scatter_(1, label.view(-1, 1).long(), 1)
output = (one_hot * phi) + ((1.0 - one_hot) * cosine)
output *= self.s
return output
if __name__ == "__main__":
# input = Variable(torch.FloatTensor(2, 3, 112, 96))
input = Variable(torch.FloatTensor(2, 3, 128, 128)) # 해상도 128x128 수정 진행.
net = MobileFacenet()
print(net)
x = net(input)
print(x.shape)

19
core/utils.py
Unescape View File

@ -0,0 +1,19 @@
from __future__ import print_function
import os
import logging
def init_log(output_dir):
logging.basicConfig(level=logging.DEBUG,
format='%(asctime)s %(message)s',
datefmt='%Y%m%d-%H:%M:%S',
filename=os.path.join(output_dir, 'log.log'),
filemode='w')
console = logging.StreamHandler()
console.setLevel(logging.INFO)
logging.getLogger('').addHandler(console)
return logging
if __name__ == '__main__':
pass

50
dataloader/CASIA_Face_loader.py
Unescape View File

@ -0,0 +1,50 @@
import numpy as np
import scipy.misc
import os
import torch
class CASIA_Face(object):
def __init__(self, root):
self.root = root
img_txt_dir = os.path.join(root, 'CASIA-WebFace-112X96.txt')
image_list = []
label_list = []
with open(img_txt_dir) as f:
img_label_list = f.read().splitlines()
for info in img_label_list:
image_dir, label_name = info.split(' ')
image_list.append(os.path.join(root, 'CASIA-WebFace-112X96', image_dir))
label_list.append(int(label_name))
self.image_list = image_list
self.label_list = label_list
self.class_nums = len(np.unique(self.label_list))
def __getitem__(self, index):
img_path = self.image_list[index]
target = self.label_list[index]
img = scipy.misc.imread(img_path)
if len(img.shape) == 2:
img = np.stack([img] * 3, 2)
flip = np.random.choice(2)*2-1
img = img[:, ::flip, :]
img = (img - 127.5) / 128.0
img = img.transpose(2, 0, 1)
img = torch.from_numpy(img).float()
return img, target
def __len__(self):
return len(self.image_list)
if __name__ == '__main__':
data_dir = '/home/brl/USER/fzc/dataset/CASIA'
dataset = CASIA_Face(root=data_dir)
trainloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True, num_workers=8, drop_last=False)
print(len(dataset))
for data in trainloader:
print(data[0].shape)

33
dataloader/LFW_loader.py
Unescape View File

@ -0,0 +1,33 @@
import numpy as np
import scipy.misc
import torch
class LFW(object):
def __init__(self, imgl, imgr):
self.imgl_list = imgl
self.imgr_list = imgr
def __getitem__(self, index):
imgl = scipy.misc.imread(self.imgl_list[index])
if len(imgl.shape) == 2:
imgl = np.stack([imgl] * 3, 2)
imgr = scipy.misc.imread(self.imgr_list[index])
if len(imgr.shape) == 2:
imgr = np.stack([imgr] * 3, 2)
# imgl = imgl[:, :, ::-1]
# imgr = imgr[:, :, ::-1]
imglist = [imgl, imgl[:, ::-1, :], imgr, imgr[:, ::-1, :]]
for i in range(len(imglist)):
imglist[i] = (imglist[i] - 127.5) / 128.0
imglist[i] = imglist[i].transpose(2, 0, 1)
imgs = [torch.from_numpy(i).float() for i in imglist]
return imgs
def __len__(self):
return len(self.imgl_list)
if __name__ == '__main__':
pass

95
dataloader/MyHF_loader.py
Unescape View File

@ -0,0 +1,95 @@
import torch
from torch.utils.data import Dataset
from PIL import Image
import numpy as np
from datasets import load_dataset
# ----------------------------
# Train Dataset: CASIA Web Face
# ----------------------------
class CASIA_HF(Dataset):
def __init__(self):
self.dataset = load_dataset("SaffalPoosh/casia_web_face", split="train") # Hugging Face train split
def __len__(self):
return len(self.dataset)
def __getitem__(self, idx):
item = self.dataset[idx]
img = np.array(item['image']) # Hugging Face image 열
img = Image.fromarray(img).convert("RGB").resize((128,128))
img = np.array(img)
img = (img - 127.5) / 128.0
img = img.transpose(2,0,1)
img = torch.from_numpy(img).float()
label = torch.tensor(int(item['label'])) # label 열 확인 필요
return img, label
# ----------------------------
# Test Dataset: LFW Pairs
# ----------------------------
# class LFW_Pairs(Dataset):
# def __init__(self):
# self.dataset = load_dataset("logasja/lfw", "pairs", split="test")
# def __len__(self):
# return len(self.dataset)
# def __getitem__(self, idx):
# item = self.dataset[idx]
# imgl = np.array(item['image1'])
# imgr = np.array(item['image2'])
# imgl = Image.fromarray(imgl).convert("RGB").resize((128,128))
# imgr = Image.fromarray(imgr).convert("RGB").resize((128,128))
# imglist = [imgl, imgl[:, ::-1, :], imgr, imgr[:, ::-1, :]] # original + flip
# for i in range(len(imglist)):
# imglist[i] = (imglist[i] - 127.5) / 128.0
# imglist[i] = imglist[i].transpose(2,0,1)
# imgs = [torch.from_numpy(i).float() for i in imglist]
# label = torch.tensor(item['label'])
# return imgs, label
class LFW_Pairs(Dataset):
def __init__(self):
self.dataset = load_dataset("logasja/lfw", "pairs", split="test")
def __len__(self):
return len(self.dataset)
def __getitem__(self, idx):
item = self.dataset[idx]
# print(idx,item) # 지울거
# print(type(item)) # 지울거
# imgl = np.array(item['img_0'])
# imgr = np.array(item['img_1'])
# PIL 이미지 가져오기
imgl = item['img_0']
imgr = item['img_1']
imgl = imgl.resize((128,128)).convert("RGB")
imgr = imgr.resize((128,128)).convert("RGB")
# print('imgl shape:', imgl.shape, 'type:', type(imgl))
# print('imgr shape:', imgr.shape, 'type:', type(imgr))
# numpy 배열로 변환
imgl = np.array(imgl)
imgr = np.array(imgr)
# print('numpy 배열로 변환 후, imgl shape:', imgl.shape, 'type:', type(imgl))
# print('numpy 배열로 변환 후, imgr shape:', imgr.shape, 'type:', type(imgr))
# imglist = [imgl, imgl[:, ::-1, :], imgr, imgr[:, ::-1, :]] # original + flip
# 이미지 리스트 생성 (original + flip)
imglist = [imgl, imgl[:, ::-1, :], imgr, imgr[:, ::-1, :]]
for i in range(len(imglist)):
imglist[i] = (imglist[i] - 127.5) / 128.0
imglist[i] = imglist[i].transpose(2, 0, 1)
imgs = [torch.from_numpy(i).float() for i in imglist]
label = torch.tensor(item['pair'])
return imgs, label

187
lfw_eval.py
Unescape View File

@ -0,0 +1,187 @@
import sys
# import caffe
import os
import numpy as np
import cv2
import scipy.io
import copy
import core.model
import os
import torch.utils.data
from core import model
from dataloader.LFW_loader import LFW
from config import LFW_DATA_DIR
import argparse
def parseList(root):
with open(os.path.join(root, 'pairs.txt')) as f:
pairs = f.read().splitlines()[1:]
folder_name = 'lfw-112X96'
nameLs = []
nameRs = []
folds = []
flags = []
for i, p in enumerate(pairs):
p = p.split('\t')
if len(p) == 3:
nameL = os.path.join(root, folder_name, p[0], p[0] + '_' + '{:04}.jpg'.format(int(p[1])))
nameR = os.path.join(root, folder_name, p[0], p[0] + '_' + '{:04}.jpg'.format(int(p[2])))
fold = i // 600
flag = 1
elif len(p) == 4:
nameL = os.path.join(root, folder_name, p[0], p[0] + '_' + '{:04}.jpg'.format(int(p[1])))
nameR = os.path.join(root, folder_name, p[2], p[2] + '_' + '{:04}.jpg'.format(int(p[3])))
fold = i // 600
flag = -1
nameLs.append(nameL)
nameRs.append(nameR)
folds.append(fold)
flags.append(flag)
# print(nameLs)
return [nameLs, nameRs, folds, flags]
def getAccuracy(scores, flags, threshold):
p = np.sum(scores[flags == 1] > threshold)
n = np.sum(scores[flags == -1] < threshold)
return 1.0 * (p + n) / len(scores)
def getThreshold(scores, flags, thrNum):
accuracys = np.zeros((2 * thrNum + 1, 1))
thresholds = np.arange(-thrNum, thrNum + 1) * 1.0 / thrNum
for i in range(2 * thrNum + 1):
accuracys[i] = getAccuracy(scores, flags, thresholds[i])
max_index = np.squeeze(accuracys == np.max(accuracys))
bestThreshold = np.mean(thresholds[max_index])
return bestThreshold
def evaluation_10_fold(root='./result/pytorch_result.mat'):
ACCs = np.zeros(10)
result = scipy.io.loadmat(root)
for i in range(10):
fold = result['fold']
flags = result['flag']
featureLs = result['fl']
featureRs = result['fr']
valFold = fold != i
testFold = fold == i
flags = np.squeeze(flags)
mu = np.mean(np.concatenate((featureLs[valFold[0], :], featureRs[valFold[0], :]), 0), 0)
mu = np.expand_dims(mu, 0)
featureLs = featureLs - mu
featureRs = featureRs - mu
featureLs = featureLs / np.expand_dims(np.sqrt(np.sum(np.power(featureLs, 2), 1)), 1)
featureRs = featureRs / np.expand_dims(np.sqrt(np.sum(np.power(featureRs, 2), 1)), 1)
scores = np.sum(np.multiply(featureLs, featureRs), 1)
threshold = getThreshold(scores[valFold[0]], flags[valFold[0]], 10000)
ACCs[i] = getAccuracy(scores[testFold[0]], flags[testFold[0]], threshold)
# print('{} {:.2f}'.format(i+1, ACCs[i] * 100))
# print('--------')
# print('AVE {:.2f}'.format(np.mean(ACCs) * 100))
return ACCs
def getFeatureFromTorch(lfw_dir, feature_save_dir, resume=None, gpu=True):
net = model.MobileFacenet()
if gpu:
net = net.cuda()
if resume:
ckpt = torch.load(resume)
net.load_state_dict(ckpt['net_state_dict'])
net.eval()
nl, nr, flods, flags = parseList(lfw_dir)
lfw_dataset = LFW(nl, nr)
lfw_loader = torch.utils.data.DataLoader(lfw_dataset, batch_size=32,
shuffle=False, num_workers=8, drop_last=False)
featureLs = None
featureRs = None
count = 0
for data in lfw_loader:
if gpu:
for i in range(len(data)):
data[i] = data[i].cuda()
count += data[0].size(0)
print('extracing deep features from the face pair {}...'.format(count))
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)
# featureLs.append(featureL)
# featureRs.append(featureR)
result = {'fl': featureLs, 'fr': featureRs, 'fold': flods, 'flag': flags}
scipy.io.savemat(feature_save_dir, result)
# def getFeatureFromCaffe(gpu=True):
# if gpu:
# caffe.set_mode_gpu()
# caffe.set_device(0)
# else:
# caffe.set_mode_cpu()
# # caffe.reset_all()
# model = '/home/xiaocc/Documents/caffe_project/sphereface/train/code/sphereface_deploy.prototxt'
# weights = '/home/xiaocc/Documents/caffe_project/sphereface/train/result/sphereface_model.caffemodel'
# net = caffe.Net(model, weights, caffe.TEST)
#
# nl, nr, flods, flags = parseList()
#
# featureLs = []
# featureRs = []
# for i in range(len(nl)):
# print('extracing deep features from the {}th face pair ...'.format(i))
# featureL = extractDeepFeature(nl[i], net)[0]
# featureR = extractDeepFeature(nr[i], net)[0]
# featureLs.append(featureL)
# featureRs.append(featureR)
# result = {'fl': featureLs, 'fr': featureRs, 'fold': flods, 'flag': flags}
# scipy.io.savemat('caffe_result.mat', result)
#
# def extractDeepFeature(f, net, h=112, w=96):
# img = cv2.imread(f)
# img = (img - 127.5) / 128
# img = img.transpose((2, 0, 1))
# net.blobs['data'].reshape(1, 3, h, w)
# net.blobs['data'].data[0, ...] = img
# res = copy.deepcopy(net.forward()['fc5'])
# net.blobs['data'].data[0, ...] = img[:, :, ::-1]
# res_ = copy.deepcopy(net.forward()['fc5'])
# r = np.concatenate((res, res_), 1)
# return r
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Testing')
parser.add_argument('--lfw_dir', type=str, default=LFW_DATA_DIR, help='The path of lfw data')
parser.add_argument('--resume', type=str, default='./model/best/068.ckpt',
help='The path pf save model')
parser.add_argument('--feature_save_dir', type=str, default='./result/best_result.mat',
help='The path of the extract features save, must be .mat file')
args = parser.parse_args()
# getFeatureFromCaffe()
getFeatureFromTorch(args.lfw_dir, args.feature_save_dir, args.resume)
ACCs = evaluation_10_fold(args.feature_save_dir)
for i in range(len(ACCs)):
print('{} {:.2f}'.format(i+1, ACCs[i] * 100))
print('--------')
print('AVE {:.2f}'.format(np.mean(ACCs) * 100))

BIN
result/best_result.mat
View File

Binary file not shown.

BIN
result/tmp_result.mat
View File

Binary file not shown.

179
test.ipynb
Unescape View File

@ -0,0 +1,179 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": 3,
"id": "2c786740",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/home/cuuva/anaconda3/envs/mfn/lib/python3.8/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
" from .autonotebook import tqdm as notebook_tqdm\n",
"Generating train split: 100%|██████████| 1000/1000 [00:00<00:00, 56341.73 examples/s]\n",
"Generating test split: 100%|██████████| 2200/2200 [00:00<00:00, 96950.62 examples/s]\n"
]
}
],
"source": [
"from datasets import load_dataset\n",
"\n",
"ds = load_dataset(\"logasja/lfw\", \"pairs\")"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "cf343f72",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"Generating train split: 100%|██████████| 13233/13233 [00:00<00:00, 29211.85 examples/s]\n"
]
}
],
"source": [
"from datasets import load_dataset\n",
"\n",
"ds = load_dataset(\"logasja/lfw\", \"default\")"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "14ee413a",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Saving train images...\n"
]
},
{
"ename": "KeyError",
"evalue": "'image1'",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mKeyError\u001b[0m Traceback (most recent call last)",
"Cell \u001b[0;32mIn[5], line 29\u001b[0m\n\u001b[1;32m 27\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mSaving train images...\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m 28\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i, item \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(train_data):\n\u001b[0;32m---> 29\u001b[0m img1 \u001b[38;5;241m=\u001b[39m Image\u001b[38;5;241m.\u001b[39mfromarray(\u001b[43mitem\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mimage1\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m]\u001b[49m)\u001b[38;5;241m.\u001b[39mconvert(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mRGB\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m 30\u001b[0m img2 \u001b[38;5;241m=\u001b[39m Image\u001b[38;5;241m.\u001b[39mfromarray(item[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mimage2\u001b[39m\u001b[38;5;124m'\u001b[39m])\u001b[38;5;241m.\u001b[39mconvert(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mRGB\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m 31\u001b[0m label \u001b[38;5;241m=\u001b[39m item[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mlabel\u001b[39m\u001b[38;5;124m'\u001b[39m] \u001b[38;5;66;03m# 0: same, 1: different\u001b[39;00m\n",
"\u001b[0;31mKeyError\u001b[0m: 'image1'"
]
}
],
"source": [
"from datasets import load_dataset\n",
"import os\n",
"from PIL import Image\n",
"import numpy as np\n",
"\n",
"# ----------------------------\n",
"# 경로 설정\n",
"# ----------------------------\n",
"LOCAL_DATA_DIR = \"/home/cuuva/lfw_images\" # 저장할 최상위 폴더\n",
"TRAIN_DIR = os.path.join(LOCAL_DATA_DIR, \"train\")\n",
"TEST_DIR = os.path.join(LOCAL_DATA_DIR, \"test\")\n",
"\n",
"os.makedirs(TRAIN_DIR, exist_ok=True)\n",
"os.makedirs(TEST_DIR, exist_ok=True)\n",
"\n",
"# ----------------------------\n",
"# Hugging Face LFW 불러오기\n",
"# ----------------------------\n",
"dataset = load_dataset(\"logasja/lfw\", \"pairs\")\n",
"\n",
"train_data = dataset[\"train\"]\n",
"test_data = dataset[\"test\"]\n",
"\n",
"# ----------------------------\n",
"# train 데이터 저장\n",
"# ----------------------------\n",
"print(\"Saving train images...\")\n",
"for i, item in enumerate(train_data):\n",
" img1 = Image.fromarray(item['image1']).convert(\"RGB\")\n",
" img2 = Image.fromarray(item['image2']).convert(\"RGB\")\n",
" label = item['label'] # 0: same, 1: different\n",
"\n",
" # 파일 이름 예: train_00001_1.jpg, train_00001_2.jpg\n",
" img1.save(os.path.join(TRAIN_DIR, f\"train_{i}_1.jpg\"))\n",
" img2.save(os.path.join(TRAIN_DIR, f\"train_{i}_2.jpg\"))\n",
"\n",
"print(f\"Train images saved: {len(train_data)*2}\")\n",
"\n",
"# ----------------------------\n",
"# test 데이터 저장\n",
"# ----------------------------\n",
"print(\"Saving test images...\")\n",
"for i, item in enumerate(test_data):\n",
" img1 = Image.fromarray(item['image1']).convert(\"RGB\")\n",
" img2 = Image.fromarray(item['image2']).convert(\"RGB\")\n",
" label = item['label']\n",
"\n",
" # 파일 이름 예: test_00001_1.jpg, test_00001_2.jpg\n",
" img1.save(os.path.join(TEST_DIR, f\"test_{i}_1.jpg\"))\n",
" img2.save(os.path.join(TEST_DIR, f\"test_{i}_2.jpg\"))\n",
"\n",
"print(f\"Test images saved: {len(test_data)*2}\")\n"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "b5830c3f",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"['pair', 'img_0', 'img_1']\n",
"{'pair': 1, 'img_0': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=250x250 at 0x7C0069CF23A0>, 'img_1': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=250x250 at 0x7C0069CF27F0>}\n"
]
}
],
"source": [
"from datasets import load_dataset\n",
"\n",
"ds = load_dataset(\"logasja/lfw\", \"pairs\", split=\"test\")\n",
"print(ds.column_names) # 현재 컬럼 이름 확인\n",
"print(ds[0]) # 첫 번째 데이터 샘플 확인\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7e53a5b1",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "mfn",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.12"
}
},
"nbformat": 4,
"nbformat_minor": 5
}

87
toonnx.py
Unescape View File

@ -0,0 +1,87 @@
import torch
import os
from core import model # 학습할 때 썼던 model 파일을 불러와야 합니다.
# ----------------------------
# 1. 설정 (경로 및 입력 사이즈)
# ----------------------------
# 사용자님이 알려주신 ckpt 경로
# ckpt_path = '/home/cuuva/face_exp/MobileFaceNet_Pytorch/model/CASIA_B512_v2_20251124_175829/best_model/best_104.ckpt'
ckpt_path = '/home/cuuva/face_exp/MobileFaceNet_Pytorch/model/CASIA_B512_v2_20251126_173236/best_model/best_063.ckpt'
onnx_path = 'best_104.onnx' # 저장될 파일 이름
# [중요] 학습할 때 사용한 이미지 해상도와 일치해야 합니다.
# 아까 코드에서 128x128로 수정하신 것을 확인했으므로 128로 설정합니다.
input_size = (1, 3, 128, 128)
def convert():
print(f"Loading checkpoint from: {ckpt_path}")
# ----------------------------
# 2. 모델 구조 정의
# ----------------------------
# 학습 코드와 동일한 모델 클래스를 인스턴스화 합니다.
net = model.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()
# ----------------------------
# 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
#opset_version=11 # ONNX 버전 (보통 11이나 12가 호환성이 좋음)
# batch size를 가변적으로 쓰고 싶다면 아래 dynamic_axes 사용 (고정하려면 주석 처리)
#dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
)
# torch.onnx.export(
# net,
# dummy_input,
# onnx_path,
# verbose=True,
# input_names=['input'],
# output_names=['output'],
# do_constant_folding=True, # 고정 상수 연산 미리 계산
# use_external_data_format=False # external .data 파일 없이 export
# )
print(f"Success! Model saved to: {os.path.abspath(onnx_path)}")
if __name__ == "__main__":
convert()

87
toonnx2.py
Unescape View File

@ -0,0 +1,87 @@
import torch
import os
from core import model2 # 학습할 때 썼던 model 파일을 불러와야 합니다.
# ----------------------------
# 1. 설정 (경로 및 입력 사이즈)
# ----------------------------
# 사용자님이 알려주신 ckpt 경로
ckpt_path = '/home/cuuva/face_exp/MobileFaceNet_Pytorch/model/MODEL_2_20251127_174006/best_model/best_004.ckpt'
onnx_path = 'best_104.onnx' # 저장될 파일 이름
# [중요] 학습할 때 사용한 이미지 해상도와 일치해야 합니다.
# 아까 코드에서 128x128로 수정하신 것을 확인했으므로 128로 설정합니다.
input_size = (1, 3, 128, 128)
def convert():
print(f"Loading checkpoint from: {ckpt_path}")
# ----------------------------
# 2. 모델 구조 정의
# ----------------------------
# 학습 코드와 동일한 모델 클래스를 인스턴스화 합니다.
net = model2.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/model'
# 최종 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_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()

87
toonnx_bak.py
Unescape View File

@ -0,0 +1,87 @@
import torch
import os
from core import model_bak # 학습할 때 썼던 model 파일을 불러와야 합니다.
# ----------------------------
# 1. 설정 (경로 및 입력 사이즈)
# ----------------------------
# 사용자님이 알려주신 ckpt 경로
# ckpt_path = '/home/cuuva/face_exp/MobileFaceNet_Pytorch/model/CASIA_B512_v2_20251124_175829/best_model/best_104.ckpt'
ckpt_path = '/home/cuuva/face_exp/MobileFaceNet_Pytorch/model/MODEL_BAK20251127_171730/best_model/best_001.ckpt'
onnx_path = 'best_104.onnx' # 저장될 파일 이름
# [중요] 학습할 때 사용한 이미지 해상도와 일치해야 합니다.
# 아까 코드에서 128x128로 수정하신 것을 확인했으므로 128로 설정합니다.
input_size = (1, 3, 128, 128)
def convert():
print(f"Loading checkpoint from: {ckpt_path}")
# ----------------------------
# 2. 모델 구조 정의
# ----------------------------
# 학습 코드와 동일한 모델 클래스를 인스턴스화 합니다.
net = model_bak.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()
# ----------------------------
# 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
#opset_version=11 # ONNX 버전 (보통 11이나 12가 호환성이 좋음)
# batch size를 가변적으로 쓰고 싶다면 아래 dynamic_axes 사용 (고정하려면 주석 처리)
#dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
)
# torch.onnx.export(
# net,
# dummy_input,
# onnx_path,
# verbose=True,
# input_names=['input'],
# output_names=['output'],
# do_constant_folding=True, # 고정 상수 연산 미리 계산
# use_external_data_format=False # external .data 파일 없이 export
# )
print(f"Success! Model saved to: {os.path.abspath(onnx_path)}")
if __name__ == "__main__":
convert()

166
train copy.py
Unescape View File

@ -0,0 +1,166 @@
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')

202
train.py
Unescape View File

@ -0,0 +1,202 @@
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
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(SAVE_DIR, MODEL_PRE + 'v2_' + 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.MobileFacenet()
ArcMargin = model.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, net.linear1.parameters())) + 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': net.linear1.parameters(), 'weight_decay': 4e-4},
{'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")

202
train2.py
Unescape View File

@ -0,0 +1,202 @@
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 model2
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(SAVE_DIR, '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 = model2.MobileFacenet()
ArcMargin = model2.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, net.linear1.parameters())) + 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': net.linear1.parameters(), 'weight_decay': 4e-4},
{'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")

202
train_bak.py
Unescape View File

@ -0,0 +1,202 @@
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_bak
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(SAVE_DIR, 'MODEL_BAK' + 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_bak.MobileFacenet()
ArcMargin = model_bak.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, net.linear1.parameters())) + 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': net.linear1.parameters(), 'weight_decay': 4e-4},
{'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")
Write Preview
Loading…
Cancel
Save