add test mode

README.md

@@ -68,6 +68,7 @@ python test_image.py
 
 optional arguments:
 --upscale_factor              super resolution upscale factor [default value is 4]
+--test_mode                   using GPU or CPU [default value is 'GPU'](choices:['GPU', 'CPU'])
 --image_name                  test low resolution image name
 --model_name                  generator model epoch name [default value is netG_epoch_4_100.pth]
 ```

test_image.py

@@ -10,24 +10,26 @@
 
 parser = argparse.ArgumentParser(description='Test Single Image')
 parser.add_argument('--upscale_factor', default=4, type=int, help='super resolution upscale factor')
+parser.add_argument('--test_mode', default='GPU', type=str, choices=['GPU', 'CPU'], help='using GPU or CPU')
 parser.add_argument('--image_name', type=str, help='test low resolution image name')
 parser.add_argument('--model_name', default='netG_epoch_4_100.pth', type=str, help='generator model epoch name')
 opt = parser.parse_args()
 
 UPSCALE_FACTOR = opt.upscale_factor
+TEST_MODE = True if opt.test_mode == 'GPU' else False
 IMAGE_NAME = opt.image_name
 MODEL_NAME = opt.model_name
 
 model = Generator(UPSCALE_FACTOR).eval()
-if torch.cuda.is_available():
-    model = model.cuda()
-# for cpu
-# model.load_state_dict(torch.load('epochs/' + MODEL_NAME, map_location=lambda storage, loc: storage))
-model.load_state_dict(torch.load('epochs/' + MODEL_NAME))
+if TEST_MODE:
+    model.cuda()
+    model.load_state_dict(torch.load('epochs/' + MODEL_NAME))
+else:
+    model.load_state_dict(torch.load('epochs/' + MODEL_NAME, map_location=lambda storage, loc: storage))
 
 image = Image.open(IMAGE_NAME)
 image = Variable(ToTensor()(image), volatile=True).unsqueeze(0)
-if torch.cuda.is_available():
+if TEST_MODE:
     image = image.cuda()
 
 start = time.clock()