improved classify script

classify

@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 
-from fastbook import *
-from fastai.vision import *
+from fastai.vision.all import *
 import argparse
 import sys
 import os
@@ -10,19 +9,23 @@ from contextlib import redirect_stdout
 
 import torch
 
+import multiprocessing
+nproc = multiprocessing.cpu_count()
+
 def main():
 	global args, out
 
 	out = sys.stdout
 	# out = debogrify_stdout()
 
-	parser = argparse.ArgumentParser(description='Classify images with a trained neural network.')
+	parser = argparse.ArgumentParser(description='Classify images with a trained neural network.', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
 	parser.add_argument('-v', '--verbose', action='store_true', help='debug')
 	parser.add_argument('--path', type=str, default='.', help='path to model')
 	parser.add_argument('--model', type=str, default='export.pkl', help='model')
 	parser.add_argument('--cpu', action='store_true', help='run on cpu')
-	parser.add_argument('--batch', type=int, default=100, help='batch size')
+	parser.add_argument('--batch', type=int, default=10000, help='batch size')
 	parser.add_argument('--move', type=float, help='move files with certainty above the threshold')
+	parser.add_argument('--workers', type=int, help='max number of data loader workers', default=nproc)
 #	parser.add_argument('--train', action='store_true', help='train the model')
 #	parser.add_argument('--watch', type=str, help='folder to watch for images to classify')
 #	parser.add_argument('--folder', type=str, default='.', help='folder with images to classify (or train)')
@@ -41,7 +44,9 @@ def main():
 		print(model_path, file=sys.stderr)
 
 	learn = load_learner(model_path)
-	if not args.cpu:
+	if args.cpu:
+		torch.set_num_threads(nproc)
+	else:
 		try:
 			learn.dls.to('cuda')
 		except Exception as e:
@@ -94,7 +99,7 @@ def predict_batch(learn, batch):
 	vocab = learn.dls.vocab
 	with redirect_stdout(sys.stderr):
 		# In some versions the following call writes the progress bar to stdout. We can't have that!
-		dl = learn.dls.test_dl(batch)
+		dl = learn.dls.test_dl(batch, num_workers=min(args.workers, args.batch))
 		preds_all = learn.get_preds(dl=dl)
 	for i in range(0, len(batch)):
 		filename = batch[i]