mirrors
/
pytorch-image-models
mirror of https://github.com/huggingface/pytorch-image-models.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Merge pull request #626 from cvcomar/master 

fix: the exception not using default topk argument
pull/645/head
Ross Wightman 2021-05-19 11:51:59 -07:00 committed by GitHub
parent e7f0db8664 d7e1e7144a
commit b3b503c13d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
inference.py
View File

@ -114,13 +114,13 @@ def main():
_logger.info('Predict: [{0}/{1}] Time {batch_time.val:.3f} ({batch_time.avg:.3f})'.format( _logger.info('Predict: [{0}/{1}] Time {batch_time.val:.3f} ({batch_time.avg:.3f})'.format(
batch_idx, len(loader), batch_time=batch_time)) batch_idx, len(loader), batch_time=batch_time))
topk_ids = np.concatenate(topk_ids, axis=0).squeeze() topk_ids = np.concatenate(topk_ids, axis=0)
with open(os.path.join(args.output_dir, './topk_ids.csv'), 'w') as out_file: with open(os.path.join(args.output_dir, './topk_ids.csv'), 'w') as out_file:
filenames = loader.dataset.filenames(basename=True) filenames = loader.dataset.filenames(basename=True)
for filename, label in zip(filenames, topk_ids): for filename, label in zip(filenames, topk_ids):
out_file.write('{0},{1},{2},{3},{4},{5}\n'.format( out_file.write('{0},{1}\n'.format(
filename, label[0], label[1], label[2], label[3], label[4])) filename, ','.join([ str(v) for v in label])))
if __name__ == '__main__': if __name__ == '__main__':