def test(**kwargs):
conf.parse(kwargs)
model = Network().eval()
if conf.LOAD_MODEL_PATH:
print(conf.LOAD_MODEL_PATH)
model.load_state_dict(torch.load(conf.CHECKPOINTS_ROOT + conf.LOAD_MODEL_PATH))
device = torch.device('cuda:0' if conf.USE_GPU else 'cpu')
model.to(device)
test_set = ImageFolder(conf.TEST_DATA_ROOT, transform)
test_loader = DataLoader(test_set, conf.BATCH_SIZE,
shuffle=False,
num_workers=conf.NUM_WORKERS)
results = list()
with torch.no_grad():
for step, (inputs, targets) in enumerate(test_loader):
inputs, targets = inputs.to(device), targets.to(device)
outs = model(inputs)
pred = torch.max(outs, 1)[1]
# print((targets == pred).float())
# (prob_top_k, idxs_top_k) = probability.topk(3, dim=1)
acc = (pred == targets).float().sum() / len(targets)
results += ((pred == targets).float().to('cpu').numpy().tolist())
print('[%5d] acc: %.3f' % (step + 1, acc))
results = np.array(results)
print('Top 1 acc: %.3f' % (np.sum(results) / len(results)))
def train(**kwargs):
conf.parse(kwargs)
# train_set = DataSet(cfg, train=True, test=False)
train_set = ImageFolder(conf.TRAIN_DATA_ROOT, transform)
train_loader = DataLoader(train_set, conf.BATCH_SIZE,
shuffle=True,
num_workers=conf.NUM_WORKERS)
model = Network()
if conf.LOAD_MODEL_PATH:
print(conf.LOAD_MODEL_PATH)
model.load_state_dict(torch.load(conf.CHECKPOINTS_ROOT + conf.LOAD_MODEL_PATH))
device = torch.device('cuda:0' if conf.USE_GPU else 'cpu')
criterion = nn.CrossEntropyLoss().to(device)
lr = conf.LEARNING_RATE
optim = torch.optim.Adam(params=model.parameters(),
lr=lr,
weight_decay=conf.WEIGHT_DECAY)
model.to(device)
for epoch in range(conf.MAX_EPOCH):
model.train()
running_loss = 0
for step, (inputs, targets) in tqdm(enumerate(train_loader)):
inputs, targets = inputs.to(device), targets.to(device)
optim.zero_grad()
outs = model(inputs)
loss = criterion(outs, targets)
loss.backward()
optim.step()
running_loss += loss.item()
if step % conf.PRINT_FREQ == conf.PRINT_FREQ - 1:
running_loss = running_loss / conf.PRINT_FREQ
print('[%d, %5d] loss: %.3f' % (epoch + 1, step + 1, running_loss))
# vis.plot('loss', running_loss)
running_loss = 0
torch.save(model.state_dict(), conf.CHECKPOINTS_ROOT + time.strftime('%Y-%m-%d-%H-%M-%S.pth'))
for param_group in optim.param_groups:
lr *= conf.LEARNING_RATE_DECAY
param_group['lr'] = lr
cfg.TRAIN.NMS_KSIZE,
cfg.TRAIN.TOPK,
cfg.MODEL.GAUSSIAN_KSIZE,
cfg.MODEL.GAUSSIAN_SIGMA,
cfg.MODEL.KSIZE,
cfg.MODEL.padding,
cfg.MODEL.dilation,
cfg.MODEL.scale_list,
)
des = HardNetNeiMask(cfg.HARDNET.MARGIN, cfg.MODEL.COO_THRSH)
model = Network(det, des, cfg.LOSS.SCORE, cfg.LOSS.PAIR, cfg.PATCH.SIZE,
cfg.TRAIN.TOPK)
print(f"{gct()} : to device")
device = torch.device("cuda")
model = model.to(device)
resume = args.resume
print(f"{gct()} : in {resume}")
checkpoint = torch.load(resume)
model.load_state_dict(checkpoint["state_dict"])
###############################################################################
# detect and compute
###############################################################################
img1_path, img2_path = args.imgpath.split("@")
kp1, des1, img1, _, _ = model.detectAndCompute(img1_path, device,
(600, 460))
kp2, des2, img2, _, _ = model.detectAndCompute(img2_path, device,
(460, 600)) #(460, 600)
# kp2, des2, img2,_,_ = model.detectAndCompute(img2_path, device, (600, 460)) #(460, 600)
