def main():
if not torch.cuda.is_available():
print(gct(), 'no gpu device available')
sys.exit(1)
print(gct(), 'Args = %s', args)
cfg = configparser.ConfigParser()
cfg.read('settings.conf')
if args.cpu:
device = torch.device('cpu')
use_cuda = False
else:
device = torch.device('cuda')
use_cuda = True
criterion = D2Loss_hpatches(scaling_steps=3)
criterion = criterion.to(device)
model = D2Net(model_file=args.model_path, use_cuda=use_cuda).to(device)
print(gct(), 'Param size = %fMB', count_parameters_in_MB(model))
if args.dataset[:4] == 'view':
csv_file = cfg['hpatches']['view_csv']
root_dir = cfg['hpatches'][f'{args.dataset}_root']
dataset = HPatchesDataset(
csv_file=csv_file,
root_dir=root_dir,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=json.loads(cfg['hpatches']['view_mean']),
std=json.loads(cfg['hpatches']['view_std']))
]),
test_flag=True)
print(gct(), 'Load test dataset.')
print(gct(), f'Root dir: {root_dir}. #Image pair: {len(dataset)}')
dataset_len = len(dataset)
split_idx = list(range(dataset_len))
part_pos = [
int(dataset_len * float(cfg['hpatches']['train_part_pos'])),
int(dataset_len * float(cfg['hpatches']['eval_part_pos']))
]
test_queue = DataLoader(
dataset,
batch_size=int(cfg['params']['train_bs']),
sampler=sampler.SubsetRandomSampler(
split_idx[part_pos[1]:]), # split_idx[part_pos[1]:]
shuffle=False,
pin_memory=True,
num_workers=2)
result_folder = 'results/'
if not os.path.exists(result_folder):
os.makedirs(result_folder)
test_acc = infer(test_queue, model, criterion, result_folder, device)
print('Test loss', test_acc)
import config as cfg
from lib.test_loader import TestDataSet
from tqdm import tqdm
from lib.metric import roc, f1_score_calc
import os
from lib.network import Network
import time
from lib.utils import count_parameters_in_MB
torch.cuda.set_device(cfg.CUDA_NUM)
test_set = TestDataSet(cfg)
net = Network()
print('Net contain %.2fMB' % count_parameters_in_MB(net))
net.cuda()
weight_path = os.path.join(cfg.weights_dir, 'net.pth')
net.load_state_dict(torch.load(weight_path))
# ------------------- Testing ------------------------
torch.set_grad_enabled(False)
net.eval()
pre_list = []
label_list = []
for test_img, test_label \
in tqdm(test_set.loader, total=test_set.batches, disable=False, ncols=0):
test_img = test_img.cuda()
print('weight_path:', weight_path)
torch.cuda.set_device(cfg.CUDA_NUM)
test_dataset = TestDataSet(cfg)
# ------------------------ Model -----------------------------
with open(genotype_path) as fp:
line = 'genotype = ' + fp.readline()
exec(line)
model = Network(cfg.init_channels, cfg.CLS_NUM, cfg.layers, genotype)
state_dict = torch.load(weight_path)
model.load_state_dict(state_dict)
model = model.cuda()
print("param size = %fMB", utils.count_parameters_in_MB(model))
model.drop_path_prob = cfg.drop_path_prob
# -------------------- Test ----------------------
torch.set_grad_enabled(False)
model.eval()
pre_list = []
label_list = []
for step, (img_batch, label_batch) in tqdm(enumerate(test_dataset.loader),
total=test_dataset.batches,
ncols=0):
img_batch = img_batch.cuda()
label_batch = label_batch.cuda()
# ------------------------ Model -----------------------------
criterion = nn.BCELoss()
model = Network(cfg.init_channels, cfg.CLS_NUM, cfg.layers, criterion)
model = model.cuda()
architect = Architect(model, cfg)
optimizer = torch.optim.SGD(model.parameters(), cfg.LR, cfg.momentum,
cfg.weight_decay)
# optimizer = torch.optim.Adam(model.parameters(), cfg.LR)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,
cfg.epochs,
eta_min=cfg.LR_min)
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
# ---------------------------------------------------------------
MAX_AUC = -1
for epoch_index in range(cfg.epochs):
scheduler.step()
lr = scheduler.get_lr()[0]
# ******************* training ************************
torch.set_grad_enabled(True)
model.train()
total_loss = 0
total_sample = 0
for step, (img_batch, label_batch, _, _, _,
_) in tqdm(enumerate(train_dataset.loader),
vis_writer = SummaryWriter(log_path)
else:
vis_writer = None
criterion = D2Loss_hpatches(scaling_steps=3, device=device).to(device)
model = D2Net(model_file=args.resume, use_cuda=use_cuda)
# Optimizer
optimizer = optim.Adam(
filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr
)
print(gct(), 'Param size = %fMB', count_parameters_in_MB(model))
# Resume training
if args.resume:
if os.path.isfile(args.resume):
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch'] + 1
optimizer.load_state_dict(checkpoint['optimizer'])
best_loss = checkpoint['loss']
seed = checkpoint['seed']
print(f'{gct()} Loaded checkpoint {args.resume} (epoch {start_epoch-1})')
else:
print(f'{gct()} No checkpoint found at {args.resume}')
raise IOError
else:
start_epoch = 0