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