def test(args):
# output folder
outdir = 'outdir'
if not os.path.exists(outdir):
os.makedirs(outdir)
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# dataset
if args.eval:
testset = get_dataset(args.dataset,
split='val',
mode='testval',
transform=input_transform)
elif args.test_val:
testset = get_dataset(args.dataset,
split='val',
mode='test',
transform=input_transform)
else:
testset = get_dataset(args.dataset,
split='test',
mode='test',
transform=input_transform)
# dataloader
loader_kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
test_data = data.DataLoader(testset,
batch_size=args.test_batch_size,
drop_last=False,
shuffle=False,
collate_fn=test_batchify_fn,
**loader_kwargs)
# model
pretrained = args.resume is None and args.verify is None
if args.model_zoo is not None:
model = get_model(args.model_zoo, pretrained=pretrained)
model.base_size = args.base_size
model.crop_size = args.crop_size
else:
# my
model_kwargs = {}
if args.choice_indices is not None:
assert 'alone_resnest50' in args.backbone
model_kwargs['choice_indices'] = args.choice_indices
#
model = get_segmentation_model(
args.model,
dataset=args.dataset,
backbone=args.backbone,
aux=args.aux,
se_loss=args.se_loss,
norm_layer=torch.nn.BatchNorm2d if args.acc_bn else SyncBatchNorm,
base_size=args.base_size,
crop_size=args.crop_size,
**model_kwargs)
# resuming checkpoint
if args.verify is not None and os.path.isfile(args.verify):
print("=> loading checkpoint '{}'".format(args.verify))
model.load_state_dict(torch.load(args.verify, map_location='cpu'))
elif args.resume is not None and os.path.isfile(args.resume):
checkpoint = torch.load(args.resume, map_location='cpu')
# strict=False, so that it is compatible with old pytorch saved models
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}'".format(args.resume))
elif not pretrained:
raise RuntimeError("=> no checkpoint found")
print(model)
if args.acc_bn:
from encoding.utils.precise_bn import update_bn_stats
data_kwargs = {
'transform': input_transform,
'base_size': args.base_size,
'crop_size': args.crop_size
}
trainset = get_dataset(args.dataset,
split=args.train_split,
mode='train',
**data_kwargs)
trainloader = data.DataLoader(ReturnFirstClosure(trainset),
batch_size=args.batch_size,
drop_last=True,
shuffle=True,
**loader_kwargs)
print('Reseting BN statistics')
#model.apply(reset_bn_statistics)
model.cuda()
update_bn_stats(model, trainloader)
if args.export:
torch.save(model.state_dict(), args.export + '.pth')
return
scales = [0.75, 1.0, 1.25, 1.5, 1.75, 2.0, 2.25] if args.dataset == 'citys' else \
[0.5, 0.75, 1.0, 1.25, 1.5, 1.75]#, 2.0
evaluator = MultiEvalModule(model, testset.num_class, scales=scales).cuda()
evaluator.eval()
metric = utils.SegmentationMetric(testset.num_class)
tbar = tqdm(test_data)
for i, (image, dst) in enumerate(tbar):
if args.eval:
with torch.no_grad():
predicts = evaluator.parallel_forward(image)
metric.update(dst, predicts)
pixAcc, mIoU = metric.get()
tbar.set_description('pixAcc: %.4f, mIoU: %.4f' %
(pixAcc, mIoU))
else:
with torch.no_grad():
outputs = evaluator.parallel_forward(image)
predicts = [
testset.make_pred(torch.max(output, 1)[1].cpu().numpy())
for output in outputs
]
for predict, impath in zip(predicts, dst):
mask = utils.get_mask_pallete(predict, args.dataset)
outname = os.path.splitext(impath)[0] + '.png'
mask.save(os.path.join(outdir, outname))
if args.eval:
print('pixAcc: %.4f, mIoU: %.4f' % (pixAcc, mIoU))
df = pd.DataFrame(predictions, columns=['ImageId_ClassId', 'EncodedPixels'])
df.to_csv(os.path.join(args.result, arch + "_submission.csv"), index=False)
# end of for
=======
if args.acc_bn:
from encoding.utils.precise_bn import update_bn_stats
data_kwargs = {'transform': input_transform, 'base_size': args.base_size,
'crop_size': args.crop_size}
trainset = get_dataset(args.dataset, split=args.train_split, mode='train', **data_kwargs)
trainloader = data.DataLoader(ReturnFirstClosure(trainset), batch_size=args.batch_size,
drop_last=True, shuffle=True, **loader_kwargs)
print('Reseting BN statistics')
#model.apply(reset_bn_statistics)
model.cuda()
update_bn_stats(model, trainloader)
if args.export:
torch.save(model.state_dict(), args.export + '.pth')
return
scales = [0.75, 1.0, 1.25, 1.5, 1.75, 2.0, 2.25] if args.dataset == 'citys' else \
[0.5, 0.75, 1.0, 1.25, 1.5, 1.75]#, 2.0
evaluator = MultiEvalModule(model, testset.num_class, scales=scales).cuda()
evaluator.eval()
metric = utils.SegmentationMetric(testset.num_class)
tbar = tqdm(test_data)
for i, (image, dst) in enumerate(tbar):
if args.eval:
with torch.no_grad():