cleaned classification_test.py a bit!

Author: Coderx7

ImageNet/training_scripts/imagenet_training/classification_test.py

@@ -13,6 +13,7 @@
 parser.add_argument('--model', '-m', metavar='MODEL', default='simpnet', help='model architecture (default: simpnet)')
 parser.add_argument('--num-classes', type=int, default=1000, help='Number classes in dataset')
 parser.add_argument('--weights', default='', type=str, metavar='PATH', help='path to model weights (default: none)')
+parser.add_argument('--pretrained', dest='pretrained', action='store_true', help='use pre-trained model')
 parser.add_argument('--jit', action='store_true', default=False, help='convert the model to jit before doing classification!')
 parser.add_argument('--netscale', type=float, default=1.0, help='scale of the net (default 1.0)')
 parser.add_argument('--netidx', type=int, default=0, help='which network to use (5mil or 8mil)')
@@ -25,23 +26,22 @@
 model = create_model(
     args.model,
     num_classes=args.num_classes,
+    pretrained=args.pretrained,
     checkpoint_path=args.weights,
     scale=args.netscale,
     network_idx = args.netidx,
     mode = args.netmode,
     )
+model.eval()
 
-# print('Restoring model state from checkpoint...')
-# model_weights = torch.load(args.weights, map_location='cpu')
-# model.load_state_dict(model_weights)
-# model.eval()
+if not args.pretrained and not args.weights:
+    print(f'WARNING: No pretrained weights specified! (pretrained is False and there is no checkpoint specified!)') 
 
 if args.jit:
     dummy_input = torch.randn(1, 3, 224, 224, device="cpu")
     model = torch.jit.trace(model, dummy_input)
 
 config = resolve_data_config({}, model=model)
-print(f'config: {config}')
 transform = create_transform(**config)
 
 filename = "./misc_files/dog.jpg"
@@ -53,13 +53,14 @@
 with torch.no_grad():
     out = model(tensor)
 probabilities = torch.nn.functional.softmax(out[0], dim=0)
-print(probabilities.shape) # prints: torch.Size([1000])
+print(f'{probabilities.shape}') # prints: torch.Size([1000])
 
 filename="./misc_files/imagenet_classes.txt"
 with open(filename, "r") as f:
     categories = [s.strip() for s in f.readlines()]
 
 # Print top categories per image
+print(f'Top categories:')
 top5_prob, top5_catid = torch.topk(probabilities, 5)
 for i in range(top5_prob.size(0)):
     print(categories[top5_catid[i]], top5_prob[i].item())