def eval(checkpoint, data_path, params):
# 数据
files, label, boxes = load_annotation(data_path, 'test')
eval_set = YoloDataset(paths=files,
bboxes=boxes,
labels=label,
params=params,
train=False)
eval_loader = DataLoader(eval_set,
batch_size=params.batch_size,
num_workers=params.num_gpus * 8,
shuffle=False)
# 模型
state_dict = torch.load(checkpoint)
model = Backbone()
model.load_state_dict(state_dict)
model = model.cuda()
# 损失
criterion = SumSquareError()
model.eval()
total_loss = 0
with torch.no_grad():
for iter, (img, annotation) in enumerate(eval_loader):
img = img.cuda()
annotation = annotation.cuda()
output = model(img)
loss = criterion(output, annotation).item()
total_loss += loss * len(img)
print(f'evaluate loss: {total_loss / len(eval_set)}')
def get_embeddings(data_root,
model_root,
input_size=[112, 112],
embedding_size=512):
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# check data and model paths
assert os.path.exists(data_root)
assert os.path.exists(model_root)
print(f"Data root: {data_root}")
# define image preprocessing
transform = transforms.Compose(
[
transforms.Resize([
int(128 * input_size[0] / 112),
int(128 * input_size[0] / 112)
], ), # smaller side resized
transforms.CenterCrop([input_size[0], input_size[1]]),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
], )
# define data loader
dataset = datasets.ImageFolder(data_root, transform)
loader = data.DataLoader(
dataset,
batch_size=1,
shuffle=False,
pin_memory=True,
num_workers=0,
)
print(f"Number of classes: {len(loader.dataset.classes)}")
# load backbone weigths from a checkpoint
backbone = Backbone(input_size)
backbone.load_state_dict(
torch.load(model_root, map_location=torch.device("cpu")))
backbone.to(device)
backbone.eval()
# get embedding for each face
embeddings = np.zeros([len(loader.dataset), embedding_size])
with torch.no_grad():
for idx, (image, _) in enumerate(
tqdm(loader,
desc="Create embeddings matrix",
total=len(loader)), ):
embeddings[idx, :] = F.normalize(backbone(image.to(device))).cpu()
# get all original images
images = []
for img_path, _ in dataset.samples:
img = cv2.imread(img_path)
images.append(img)
return images, embeddings
def train(params, _run=None):
params = Params(params)
set_random_seeds(params.seed)
time_now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
params.save_root = params.save_root + f'/{params.project_name}_{time_now}_{params.version}'
os.makedirs(params.save_root, exist_ok=True)
logging.basicConfig(filename=f'{params.save_root}/{params.project_name}_{time_now}_{params.version}.log',
filemode='a', format='%{asctime}s - %(levalname)s: %(message)s')
if params.num_gpus == 0:
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
logging.info(f'Available GPUs: {torch.cuda.device_count()}')
train2007, train_label_2007, train_bb_2007 = load_annotation(os.path.join(params.data_root, 'VOC2007'), 'trainval')
test2007, test_label_2007, test_bb_2007 = load_annotation(os.path.join(params.data_root, 'VOC2007'), 'test')
train2012, train_label_2012, train_bb_2012 = load_annotation(os.path.join(params.data_root, 'VOC2012'), 'trainval')
test2012, test_label_2012, test_bb_2012 = load_annotation(os.path.join(params.data_root, 'VOC2012'), 'test')
train_data = train2007+test2007+train2012
train_label = train_label_2007+test_label_2007+train_label_2012
train_bb = train_bb_2007 + test_bb_2007 + train_bb_2012
test_data = test2012
test_label = test_label_2012
test_bb = test_bb_2012
train_dataset = YoloDataset(train_data, train_bb, train_label, params, train=True)
eval_dataset = YoloDataset(test_data, test_bb, test_label, params, train=False)
train_loader = DataLoader(dataset=train_dataset, num_workers=params.num_gpus*8, batch_size=params.batch_size,
shuffle=True, drop_last=True, pin_memory=True)
eval_loader = DataLoader(dataset=eval_dataset, num_workers=1, batch_size=1,
shuffle=False, pin_memory=True)
model = Backbone()
last_step = 0
last_epoch = 0
if params.num_gpus > 0:
model = model.cuda()
if params.num_gpus > 1:
model = nn.DataParallel(model)
if params.optim == 'Adam':
optimizer = torch.optim.Adam(model.parameters(), lr=params.learning_rate)
else:
optimizer = torch.optim.SGD(model.parameters(), lr=params.learning_rate, momentum=0.9, nesterov=True, weight_decay=0.0005)
criterion = SumSquareError()
schedule = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer=optimizer, factor=0.5, verbose=True, patience=10)
epoch = 0
begin_epoch = max(0, last_epoch)
step = max(0, last_step)
best_loss = 1e6
logging.info('Begin to train...')
model.train()
import cv2 as cv
try:
for epoch in range(begin_epoch, params.epoch):
for iter, (img, annotation) in enumerate(train_loader):
output = model(img.cuda())
loss = criterion(output, annotation.cuda())
optimizer.zero_grad()
loss.backward()
optimizer.step()
if iter % params.save_interval == 0:
logging.info(f'{datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")} '
f'Train Epoch: {epoch} iter: {iter} loss: {loss.item()}')
step += 1
if epoch % params.eval_interval == 0:
model.eval()
epoch_loss = 0
with torch.no_grad():
for iter, (img, annotation) in enumerate(eval_loader):
output = model(img.cuda())
loss = criterion(output, annotation.cuda()).item()
epoch_loss += loss * len(img)
loss = epoch_loss / len(eval_dataset)
logging.info(f'{datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")} '
f'Eval Epoch: {epoch} loss: {loss}')
schedule.step(loss)
if loss < best_loss:
best_loss = loss
save_checkpoint(model, f'{params.save_root}/{epoch}_{step}.pth')
model.train()
except KeyboardInterrupt:
save_checkpoint(model, f'{params.save_root}/Interrupt_{epoch}_{step}.pth')
def main(args):
model = Backbone()
if args.use_gpu:
model.load_state_dict(torch.load(args.model_path), strict=True)
else:
model.load_state_dict(torch.load(args.model_path, map_location='cpu'),
strict=True)
device = 'cuda' if args.use_gpu else 'cpu'
model = model.to(device)
# print(model)
model.eval()
batch_size = 1
x = torch.randn(batch_size,
3,
args.input_size,
args.input_size,
requires_grad=True).to(device)
torch_out = model(x)
print("파이토치 모델 실행시간 측정")
test_cnt = 5
total_time = 0.
for i in range(test_cnt):
x = torch.randn(batch_size,
3,
args.input_size,
args.input_size,
requires_grad=True).to(device)
start = time.time()
torch_out = model(x)
total_time += time.time() - start
pytorch_exec_time = total_time / test_cnt
print("파이토치 모델 실행시간 측정 완료")
os.makedirs(os.path.dirname(args.onnx_output_path), exist_ok=True)
# 모델 변환
torch.onnx.export(
model, # 실행될 모델
x, # 모델 입력값 (튜플 또는 여러 입력값들도 가능)
args.onnx_output_path, # 모델 저장 경로 (파일 또는 파일과 유사한 객체 모두 가능)
export_params=True, # 모델 파일 안에 학습된 모델 가중치를 저장할지의 여부
# opset_version=11, # 모델을 변환할 때 사용할 ONNX 버전
do_constant_folding=True, # 최적하시 상수폴딩을 사용할지의 여부
input_names=['input'], # 모델의 입력값을 가리키는 이름
output_names=['output'], # 모델의 출력값을 가리키는 이름
# )
# dynamic_axes={'input': {2: 'height', 3: 'width'}}
)
print("onnx 모델 컨버팅 완료")
print("onnx 모델 로딩")
ort_session = onnxruntime.InferenceSession(args.onnx_output_path)
print("onnx 모델 로딩 완료")
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
ort_outs = ort_session.run(None, ort_inputs)
print("onnx 모델 실행시간 측정")
total_time = 0.
for i in range(test_cnt):
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
start = time.time()
ort_outs = ort_session.run(None, ort_inputs)
total_time += time.time() - start
onnx_exec_time = total_time / test_cnt
print("파이토치 모델 실행시간 측정 완료")
np.testing.assert_allclose(to_numpy(torch_out),
ort_outs[0],
rtol=1e-03,
atol=1e-05)
print(
"Exported model has been tested with ONNXRuntime, and the result looks good!"
)
print("onnx 평균시간", onnx_exec_time, "파이토치 평균시간", pytorch_exec_time)
class DogEyesVerifier(object):
def __init__(self, input_size=112, model_path=None, pos_thr=None, neg_thr=None, drop_ratio=0.6,
net_depth=50, use_random_crop=False,
use_gray=False, use_center_crop=False, center_crop_ratio=0.8, device='CPU', use_onnx=False):
assert model_path is not None
self.pos_thr = pos_thr
self.neg_thr = neg_thr
self.input_size = input_size
self.transforms = get_test_transforms(input_size, use_random_crop=use_random_crop, use_gray=use_gray,
use_center_crop=use_center_crop, center_crop_ratio=center_crop_ratio)
self.use_onnx = use_onnx
if use_onnx:
import onnxruntime
self.model = onnxruntime.InferenceSession(model_path)
else:
self.model = Backbone(net_depth, drop_ratio, 'ir_se').to(device)
self.model.eval()
if device == 'cpu':
self.model.load_state_dict(torch.load(model_path, map_location=device), strict=True)
else:
self.model.load_state_dict(torch.load(model_path), strict=True)
self.use_cuda = device == "cuda"
def is_same(self, img1, img2, is_same_side, pos_thr=None, neg_thr=None, use_pos_low_thr=False):
if pos_thr is None:
pos_thr = self.pos_thr
if neg_thr is None:
neg_thr = self.neg_thr
if pos_thr is None and neg_thr is None:
raise ValueError("pos_thr and neg_thr are None.")
if isinstance(img1, str):
img1 = cv2.imread(img1, cv2.IMREAD_COLOR)
img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
elif isinstance(img1, JpegImageFile):
if img1.mode != "RGB":
img1 = img1.convert("RGB")
img1 = np.array(img1)
if isinstance(img2, str):
img2 = cv2.imread(img2, cv2.IMREAD_COLOR)
img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2RGB)
elif isinstance(img2, JpegImageFile):
if img2.mode != "RGB":
img2 = img2.convert("RGB")
img2 = np.array(img2)
img1 = self.transforms(image=img1)['image'].unsqueeze(0)
img2 = self.transforms(image=img2)['image'].unsqueeze(0)
if self.use_cuda:
img1 = img1.cuda()
img2 = img2.cuda()
if self.use_onnx:
input1 = {self.model.get_inputs()[0].name: to_numpy(img1)}
input2 = {self.model.get_inputs()[0].name: to_numpy(img2)}
embedding1 = self.model.run(None, input1)[0]
embedding2 = self.model.run(None, input2)[0]
else:
with torch.set_grad_enabled(False):
embedding1 = self.model(img1).cpu().data.numpy()
embedding2 = self.model(img2).cpu().data.numpy()
dist = np.sum(np.square(np.subtract(embedding1, embedding2)), 1)[0]
if is_same_side:
return dist < pos_thr
else:
if use_pos_low_thr:
return dist < neg_thr
else:
return False
