def evaluate(model, weights, dataset, datatype, split, count, shot, seed, gpu,
hist_path, seg_path):
print("evaluating {} with weights {} on {} {}-{}".format(
model, weights, datatype, dataset, split))
os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
device = torch.device('cuda:0')
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
np.random.seed(seed)
prepare_data = datatypes[datatype]
dataset = prepare_data(dataset, split, count=count, shot=shot)
loader = prepare_loader(dataset, evaluation=True)
model = prepare_model(model, dataset.num_classes,
weights=weights).to(device)
model.eval()
loss_fn = nn.CrossEntropyLoss(reduction='mean',
ignore_index=dataset.ignore_index)
total_loss = 0.
metrics = SegScorer(
len(dataset.classes)
) # n.b. this is the full no. of classes, not the no. of model outputs
with torch.no_grad():
for i, data in enumerate(loader):
inputs, target, aux = data[:-2], data[-2], data[-1]
inputs = [
inp.to(device) if not isinstance(inp, list) else
[[i_.to(device) for i_ in in_] for in_ in inp]
for inp in inputs
]
target = target.to(device)
scores = model(*inputs)
loss = loss_fn(scores, target)
total_loss += loss.item()
# segmentation evaluation
_, seg = scores.data[0].max(0)
metrics.update(
seg.to('cpu').numpy(),
target.to('cpu').numpy(), aux)
# optionally save segs
if seg_path is not None:
seg = Image.fromarray(seg.to('cpu').numpy().astype(np.uint8),
mode='P')
save_id = f"{aux['slug']}_{aux.get('cls', 'all')}_{aux.get('inst', 'all')}"
seg.save(f"{seg_path}/{save_id}.png")
print("loss {}".format(total_loss / len(dataset)))
for metric, score in metrics.score().items():
score = np.nanmean(score)
print("{:10s} {:.3f}".format(metric, score))
if hist_path is not None:
metrics.save(hist_path)
def main(exp, model, weight_iter, dataset, datatype, split, seed, gpu):
print("evaluating {} with weights {} on {}".format(model, exp, dataset))
os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
weights = f"experiments/{exp}/snapshot-iter{weight_iter}.pth"
output_dir = f"experiments/{exp}"
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
np.random.seed(seed)
results = []
data_name = dataset
model_name = model
for count in [2, 4, 8, 16, 32, 64, 128]:
print(f"{count} annotations")
prepare_data = datatypes[datatype]
dataset = prepare_data(data_name, split, count=count)
loader = prepare_loader(dataset, shuffle=False)
model = prepare_model(model_name, dataset.num_classes).cuda()
model.load_state_dict(torch.load(weights))
model.eval()
loss_fn = CrossEntropyLoss2D(size_average=True,
ignore_index=dataset.ignore_index)
total_loss = 0.
metrics = SegScorer(
len(dataset.classes)
) # n.b. this is the full no. of classes, not the no. of model outputs
for i, data in enumerate(loader):
inputs, target, aux = data[:-2], data[-2], data[-1]
inputs = [Variable(in_, volatile=True).cuda() for in_ in inputs]
target = Variable(target, volatile=True).cuda(async=True)
scores = model(*inputs)
loss = loss_fn(scores, target)
total_loss += loss.data[0]
# segmentation evaluation
_, seg = scores.data[0].max(0)
metrics.update(seg.cpu().numpy(), target.data.cpu().numpy(), aux)
print("loss {}".format(total_loss / len(dataset)))
for metric, score in metrics.score().items():
if score.size > 1:
score = np.nanmean(score)
print("{:10s} {:.3f}".format(metric, score))
np.save(f"{output_dir}/hist_iter{weight_iter}_{count}_annos.npy",
metrics.hist)
def evaluate(model, weights, dataset, datatype, split, count, shot, multi,
seed, gpu, output):
print("evaluating {} with weights {} on {} {}-{}".format(
model, weights, datatype, dataset, split))
os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
np.random.seed(seed)
prepare_data = datatypes[datatype]
dataset = prepare_data(dataset, split, count=count, shot=shot, multi=multi)
loader = prepare_loader(dataset, evaluation=True)
model = prepare_model(model, dataset.num_classes, weights=weights).cuda()
model.eval()
loss_fn = CrossEntropyLoss2D(size_average=True,
ignore_index=dataset.ignore_index)
total_loss = 0.
metrics = SegScorer(
len(dataset.classes)
) # n.b. this is the full no. of classes, not the no. of model outputs
for i, data in enumerate(loader):
inputs, target, aux = data[:-2], data[-2], data[-1]
inputs = [
Variable(inp, volatile=True).cuda() if not isinstance(inp, list)
else [[Variable(i_, volatile=True).cuda() for i_ in in_]
for in_ in inp] for inp in inputs
]
target = Variable(target, volatile=True).cuda(async=True)
scores = model(*inputs)
loss = loss_fn(scores, target)
total_loss += loss.data[0]
# segmentation evaluation
_, seg = scores.data[0].max(0)
metrics.update(seg.cpu().numpy(), target.data.cpu().numpy(), aux)
print("loss {}".format(total_loss / len(dataset)))
for metric, score in metrics.score().items():
score = np.nanmean(score)
print("{:10s} {:.3f}".format(metric, score))
if output:
metrics.save(output)
def main(experiment, model, dataset, datatype, split, val_dataset, val_split,
class_group, count, shot, lr, max_iter, seed, gpu, do_eval):
setproctitle.setproctitle(experiment)
version = subprocess.check_output(['git', 'describe', '--always'],
universal_newlines=True).strip()
# experiment metadata
args = locals()
exp_dir = './experiments/{}/'.format(experiment)
if os.path.isdir(exp_dir):
click.confirm(click.style(
"{} already exists. Do you want to "
"obliterate it and continue?".format(experiment),
fg='red'),
abort=True)
shutil.rmtree(exp_dir)
try:
os.makedirs(exp_dir, exist_ok=True)
except:
raise Exception("Could not create experiment dir {}".format(exp_dir))
os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
device = torch.device('cuda:0')
logging.basicConfig(filename=exp_dir + 'log',
level=logging.INFO,
format='[%(asctime)s.%(msecs)03d] %(message)s',
datefmt='%H:%M:%S')
logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
logging.info("training %s", experiment)
logging.info("args: %s", args)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
np.random.seed(seed)
# spawn persistent evaluation process
if do_eval:
mp_ctx = multiprocessing.get_context('spawn')
q = mp_ctx.JoinableQueue()
p = mp_ctx.Process(target=pevaluate, args=(q, ))
p.start()
# filter classes by group for heldout experiments
classes_to_filter = None
if class_group != 'all':
class_group = int(class_group)
# divide classes into quarters and take background + the given quarter
group_size = len(datasets[dataset].classes) // 4
group_idx = 1 + class_group * group_size
group_classes = range(group_idx, group_idx + group_size)
classes_to_filter = list(
set(range(1, len(datasets[dataset].classes))) - set(group_classes))
dataset_name = dataset
prepare_data = datatypes[datatype]
dataset = prepare_data(dataset_name, split, classes_to_filter, count, shot)
loader = prepare_loader(dataset)
val_dataset_name = val_dataset or dataset_name
model_name = model
model = prepare_model(model, dataset.num_classes).cuda()
loss_fn = nn.CrossEntropyLoss(reduction='mean',
ignore_index=dataset.ignore_index)
learned_params = filter(lambda p: p.requires_grad, model.parameters())
opt = optim.SGD(learned_params, lr=lr, momentum=0.99, weight_decay=0.0005)
iter_order = int(np.log10(max_iter) + 1) # for pretty printing
epoch = 0
iteration = 0
losses = []
model.train()
while iteration < max_iter:
logging.info("epoch %d", epoch)
epoch += 1
train_loss = 0.
for i, data in enumerate(loader):
inputs, target, aux = data[:-2], data[-2], data[-1]
inputs = [
inp.to(device) if not isinstance(inp, list) else
[[i_.to(device) for i_ in in_] for in_ in inp]
for inp in inputs
]
target = target.to(device, non_blocking=True)
scores = model(*inputs)
loss = loss_fn(scores, target)
loss.backward()
train_loss += loss.item()
losses.append(loss.item())
if iteration % 20 == 0:
logging.info(
"%s",
"iter {iteration:{iter_order}d} loss {mean_loss:02.5f}".
format(iteration=iteration,
iter_order=iter_order,
mean_loss=np.mean(losses)))
losses = []
if iteration % 4000 == 0:
# snapshot
logging.info("snapshotting...")
snapshot_path = exp_dir + 'snapshot-iter{iteration:0{iter_order}d}.pth'.format(
iteration=iteration, iter_order=iter_order)
torch.save(model.state_dict(), snapshot_path)
# evaluate
if do_eval:
logging.info("evaluating...")
hist_path = exp_dir + 'hist-iter{iteration:0{iter_order}d}'.format(
iteration=iteration, iter_order=iter_order)
try:
# wait for the last evalution if it's still running
q.join()
except:
pass
# carry out evaluation in independent process for determinism and speed
q.put(
(model_name, snapshot_path, val_dataset_name, datatype,
val_split, count, shot, seed, gpu, hist_path, None))
# update
opt.step()
opt.zero_grad()
iteration += 1
logging.info("%s",
"train loss = {:02.5f}".format(train_loss / len(dataset)))
# signal to evaluation process that training is done
if do_eval:
q.put(None)