perf(train): 更新类别数和类别列表

Author: zjZSTU

py/lib/models/location_dataset.py

@@ -128,16 +128,20 @@ def __len__(self):
 
 
 if __name__ == '__main__':
-    root_dir = '../../data/location_dataset/'
+    # root_dir = '../../data/location_dataset/'
+    root_dir = '../../data/VOC_dataset/'
     transform = transforms.Compose([
         transforms.ToPILImage(),
         transforms.Resize((448, 448)),
         transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
     ])
 
-    cate_list = ['cucumber', 'eggplant', 'mushroom']
-    data_set = LocationDataset(root_dir, cate_list, transform, 7, 2, 3)
+    # cate_list = ['cucumber', 'eggplant', 'mushroom']
+    # data_set = LocationDataset(root_dir, cate_list, transform, 7, 2, 3)
+    cate_list = ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable',
+                 'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']
+    data_set = LocationDataset(root_dir, cate_list, transform, 7, 2, 20)
     print(data_set)
     print(len(data_set))
 

py/lib/models/multi_part_loss.py

@@ -329,15 +329,19 @@ def load_data(data_root_dir, cate_list, S=7, B=2, C=20):
 if __name__ == '__main__':
     S = 7
     B = 2
-    C = 3
-    cate_list = ['cucumber', 'eggplant', 'mushroom']
+    # C = 3
+    # cate_list = ['cucumber', 'eggplant', 'mushroom']
+    C = 20
+    cate_list = ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable',
+                 'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']
 
-    criterion = MultiPartLoss(448, 448, S=7, B=2, C=3)
+    criterion = MultiPartLoss(448, 448, S=S, B=B, C=C)
     # preds = torch.arange(637).reshape(1, 7 * 7, 13) * 0.01
     # targets = torch.ones((1, 7 * 7, 13)) * 0.01
     # loss = criterion(preds, targets)
     # print(loss)
-    data_loader = load_data('../../data/location_dataset', cate_list, S=S, B=B, C=C)
+    # data_loader = load_data('../../data/location_dataset', cate_list, S=S, B=B, C=C)
+    data_loader = load_data('../../data/VOC_dataset', cate_list, S=S, B=B, C=C)
     model = YOLO_v1(S=S, B=B, C=C)
 
     for inputs, labels in data_loader:

py/lib/models/yolo_v1.py

@@ -86,7 +86,8 @@ def forward(self, x):
 if __name__ == '__main__':
     # data = torch.randn((1, 3, 448, 448))
     data = torch.randn((1, 3, 224, 224))
-    model = YOLO_v1(7, 2, 3)
+    # model = YOLO_v1(7, 2, 3)
+    model = YOLO_v1(7, 2, 20)
 
     outputs = model(data)
     print(outputs.shape)

py/lib/train.py

@@ -21,9 +21,12 @@
 
 S = 7
 B = 2
-C = 3
 
-cate_list = ['cucumber', 'eggplant', 'mushroom']
+# C = 3
+# cate_list = ['cucumber', 'eggplant', 'mushroom']
+C = 20
+cate_list = ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable',
+             'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']
 
 
 def load_data(data_root_dir, cate_list, S=7, B=2, C=20):