sync

public/2023/04-03_yolov8.md

@@ -973,7 +973,7 @@ os.environ["KECAM_BACKEND"] = "torch"
 import math
 import torch
 import kecam
-import pycocotools, cv2, torchvision, tqdm  # Not using here, just in case for later error
+import pycocotools, cv2, torchvision, tqdm, h5py  # Not using here, just in case for later error
 
 global_device = torch.device("cuda:0") if torch.cuda.is_available() and int(os.environ.get("CUDA_VISIBLE_DEVICES", "0")) >= 0 else torch.device("cpu")
 
@@ -1118,9 +1118,9 @@ def build_optimizer(model, name="sgd", lr=0.01, momentum=0.937, decay=5e-4):
     return optimizer
 
 dataset_path, input_shape, batch_size, epochs, initial_epoch, warmup_epochs, optimizer_name, lr = "datasets/coco_dog_cat/detections.json", 640, 16, 100, 0, 3, "sgd", 0.01
-train_dataset, _ = kecam.coco.data.init_dataset(data_path=dataset_path, batch_size=batch_size, image_size=input_shape)
+train_dataset, _, _, num_classes = kecam.coco.data.init_dataset(data_path=dataset_path, batch_size=batch_size, image_size=input_shape, with_info=True)
 image, labels = next(iter(train_dataset))
-num_classes = labels.shape[-1] - 4
+print(">>>> num_classes:", num_classes)
 
 model = kecam.yolov8.YOLOV8_N(input_shape=(3, input_shape, input_shape), classifier_activation=None, num_classes=num_classes, pretrained=None).to(global_device)
 basic_save_name = "{}_{}".format(model.name, os.path.basename(dataset_path))