if args.resume:
print('raw load')
model_raw.load_state_dict(torch.load(args.resume).state_dict())
torch.manual_seed(args.seed + args.rank)
val_dir = os.path.join(parser.dataset, '/val')
data_config = resolve_data_config(vars(args), model=model_raw, verbose=args.local_rank == 0)
num_aug_splits = 0
val_dataset = Dataset(val_dir, load_bytes=False, class_map='')
param_count = sum([m.numel() for m in model_raw.parameters()])
logging.info('Model created, param count: %d' % (param_count))
model_raw, test_time_pool = apply_test_time_pool(model_raw, data_config, args)
crop_pct = 1.0 if test_time_pool else data_config['crop_pct']
val_loader = create_loader(
val_dataset,
input_size=data_config['input_size'],
batch_size=args.batch_size,
is_training=False,
use_prefetcher=args.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
crop_pct=crop_pct,
pin_memory=args.pin_mem,
tf_preprocessing=args.tf_preprocessing)
def main():
args = parser.parse_args()
# create model
model = create_model(
args.model,
num_classes=args.num_classes,
in_chans=3,
pretrained=args.pretrained,
checkpoint_path=args.checkpoint)
print('Model %s created, param count: %d' %
(args.model, sum([m.numel() for m in model.parameters()])))
config = resolve_data_config(model, args)
model, test_time_pool = apply_test_time_pool(model, config, args)
if args.num_gpu > 1:
model = torch.nn.DataParallel(model, device_ids=list(range(args.num_gpu))).cuda()
else:
model = model.cuda()
loader = create_loader(
Dataset(args.data),
input_size=config['input_size'],
batch_size=args.batch_size,
use_prefetcher=True,
interpolation=config['interpolation'],
mean=config['mean'],
std=config['std'],
num_workers=args.workers,
crop_pct=1.0 if test_time_pool else config['crop_pct'])
model.eval()
k = min(args.topk, args.num_classes)
batch_time = AverageMeter()
end = time.time()
topk_ids = []
with torch.no_grad():
for batch_idx, (input, _) in enumerate(loader):
input = input.cuda()
labels = model(input)
topk = labels.topk(k)[1]
topk_ids.append(topk.cpu().numpy())
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if batch_idx % args.print_freq == 0:
print('Predict: [{0}/{1}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})'.format(
batch_idx, len(loader), batch_time=batch_time))
topk_ids = np.concatenate(topk_ids, axis=0).squeeze()
with open(os.path.join(args.output_dir, './topk_ids.csv'), 'w') as out_file:
filenames = loader.dataset.filenames()
for filename, label in zip(filenames, topk_ids):
filename = os.path.basename(filename)
out_file.write('{0},{1},{2},{3},{4},{5}\n'.format(
filename, label[0], label[1], label[2], label[3], label[4]))
def main():
args = parser.parse_args()
# create model
model = create_model(args.model,
num_classes=args.num_classes,
in_chans=3,
pretrained=args.pretrained,
checkpoint_path=args.checkpoint)
print('Model %s created, param count: %d' %
(args.model, sum([m.numel() for m in model.parameters()])))
data_config = resolve_data_config(model, args)
model, test_time_pool = apply_test_time_pool(model, data_config, args)
if args.num_gpu > 1:
model = torch.nn.DataParallel(model,
device_ids=list(range(
args.num_gpu))).cuda()
else:
model = model.cuda()
criterion = nn.CrossEntropyLoss().cuda()
loader = create_loader(
Dataset(args.data),
input_size=data_config['input_size'],
batch_size=args.batch_size,
use_prefetcher=True,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
crop_pct=1.0 if test_time_pool else data_config['crop_pct'])
batch_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
model.eval()
end = time.time()
with torch.no_grad():
for i, (input, target) in enumerate(loader):
target = target.cuda()
input = input.cuda()
# compute output
output = model(input)
loss = criterion(output, target)
# measure accuracy and record loss
prec1, prec5 = accuracy(output.data, target, topk=(1, 5))
losses.update(loss.item(), input.size(0))
top1.update(prec1.item(), input.size(0))
top5.update(prec5.item(), input.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % args.print_freq == 0:
print(
'Test: [{0}/{1}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f}, {rate_avg:.3f}/s) \t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
i,
len(loader),
batch_time=batch_time,
rate_avg=input.size(0) / batch_time.avg,
loss=losses,
top1=top1,
top5=top5))
print(
' * Prec@1 {top1.avg:.3f} ({top1a:.3f}) Prec@5 {top5.avg:.3f} ({top5a:.3f})'
.format(top1=top1,
top1a=100 - top1.avg,
top5=top5,
top5a=100. - top5.avg))