Skip to content

T
tongue-diagnosis
Commit b0b0ff83 authored by zhiyang.zhou's avatar zhiyang.zhou
Browse files
Options

minor change

parent 2b6936e8
Pipeline #194 canceled with stages
      Hide whitespace changes
      Inline Side-by-side
      Showing with 3 additions and 1 deletion
      train_tongue_diagnosis.py
      View file @ b0b0ff83
      ...@@ -32,7 +32,7 @@ parser.add_argument('--log_interval', type=int, default=10, metavar='N', ...@@ -32,7 +32,7 @@ parser.add_argument('--log_interval', type=int, default=10, metavar='N',
      args = parser.parse_args() args = parser.parse_args()
      NUM_CLASSES = 7 NUM_CLASSES = 8
      use_cuda = not args.no_cuda and torch.cuda.is_available() use_cuda = not args.no_cuda and torch.cuda.is_available()
      GPUID = args.gpuid GPUID = args.gpuid
      os.environ["CUDA_VISIBLE_DEVICES"] = str(GPUID) os.environ["CUDA_VISIBLE_DEVICES"] = str(GPUID)
      ...@@ -43,6 +43,8 @@ model_dir = args.model_dir ...@@ -43,6 +43,8 @@ model_dir = args.model_dir
      if not os.path.exists(model_dir): if not os.path.exists(model_dir):
      os.makedirs(model_dir) os.makedirs(model_dir)
      torch.backends.cudnn.enabled = False
      # training data and test data settings # training data and test data settings
      kwargs = {'num_workers': 4, 'pin_memory': True} if use_cuda else {} kwargs = {'num_workers': 4, 'pin_memory': True} if use_cuda else {}
      normalizer = transforms.Normalize(mean=[x / 255.0 for x in [125.3, 123.0, 113.9]], normalizer = transforms.Normalize(mean=[x / 255.0 for x in [125.3, 123.0, 113.9]],
      ......
      Markdown is supported
      0% or
      You are about to add 0 people to the discussion. Proceed with caution.
      Finish editing this message first!
      Please register or to comment