[fix]

main.py

@@ -18,7 +18,7 @@ def parse_option():
     parser = argparse.ArgumentParser('Progressive Region Enhancement Network(PRENet) for training and testing')
 
     parser.add_argument('--batchsize', default=2, type=int, help="batch size for single GPU")
-    parser.add_argument('--dataset', type=str, default='food101', help='food2k, food101, food500')
+    parser.add_argument('--dataset', type=str, default='jkyy', help='food2k, food101, food500')
     parser.add_argument('--image_path', type=str, default="E:/datasets/food101/images/", help='path to dataset')
     parser.add_argument("--train_path", type=str, default="E:/datasets/food101/meta_data/train_full.txt", help='path to training list')
     parser.add_argument("--test_path", type=str, default="E:/datasets/food101/meta_data/test_full.txt",
@@ -172,6 +172,8 @@ def main():
         NUM_CATEGORIES = 500
     elif args.dataset == "food2k":
         NUM_CATEGORIES = 2000
+    elif args.dataset == "jkyy":
+        NUM_CATEGORIES = 1788
 
 
     net = load_model('resnet50',pretrain=False,require_grad=True,num_class=NUM_CATEGORIES)