angelinager
commented
1 year ago In case you need my code for Dataset classes along with the helper function for padding (My datasets are in YOLO format):
def pad_targets(targets: np.ndarray, max_targets: int) -> np.ndarray:
padded_targets = np.zeros((max_targets, targets.shape[-1]))
padded_targets[range(len(targets))[: max_targets]] = targets[: max_targets]
padded_targets = np.ascontiguousarray(padded_targets, dtype=np.float32)
return padded_targets
class KedenCustomDataset(Dataset):
def __init__(self, data_folder):
self.image_folder = os.path.join(data_folder, "images")
self.annot_folder = os.path.join(data_folder, "labels")
# Read data files
self.images = os.listdir(self.image_folder)
# Read annotation files
self.annotations = os.listdir(self.annot_folder)
def __len__(self):
return min(len(self.images), len(self.annotations))
def __getitem__(self, i):
resize_dim = 640
pad_dim = 6
# Read image and label
image_filename = self.images[i]
annot_filename = image_filename.split('.tif')[0] + '.txt'
image_path = os.path.join(self.image_folder, image_filename)
image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
image = cv2.resize(image, (resize_dim, resize_dim))
image_tensor = torch.tensor(np.array(image, dtype=np.float32)).permute(2, 0, 1).float()
if annot_filename in self.annotations:
labels_path = os.path.join(self.annot_folder, annot_filename)
labels = np.loadtxt(labels_path, delimiter=' ')
if labels.ndim == 1:
labels = np.reshape(labels, (1, -1))
labels = pad_targets(labels, pad_dim)
else:
labels = np.zeros((pad_dim, 5))
return image_tensor, labelsI have grayscale images so I convert image to grayscale while still keeping three channels. Also, there are some images that don't have annotation (but there's just a couple of them), so I return zeros of labels matrix (I don't know whether it's a right way to handle those images though).
The Dataset class for a second dataset is very similar to the first one except that the image files are in ".png" format instead of ".tif", and in the constructor I'm also passing the number of classes from the first dataset to shift my labels in the second dataset, since labels start from 0 in both datasets.
class PidrayCustomDataset(Dataset):
def __init__(self, data_folder, keden_labels_num):
self.keden_labels_num = keden_labels_num
self.image_folder = os.path.join(data_folder, "images")
self.annot_folder = os.path.join(data_folder, "labels")
# Read data files
self.images = os.listdir(self.image_folder)
# Read annotation files
self.annotations = os.listdir(self.annot_folder)
def __len__(self):
return min(len(self.images), len(self.annotations))
def __getitem__(self, i):
resize_dim = 640
pad_dim = 6
# Read image and label
image_filename = self.images[i]
annot_filename = image_filename.split('.png')[0] + '.txt'
image_path = os.path.join(self.image_folder, image_filename)
image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
image = cv2.resize(image, (resize_dim, resize_dim))
# image = Image.open(image_path, mode='r').resize((320, 320))
image_tensor = torch.tensor(np.array(image, dtype=np.float32)).permute(2, 0, 1).float()
if annot_filename in self.annotations:
labels_path = os.path.join(self.annot_folder, annot_filename)
labels = np.loadtxt(labels_path, delimiter=' ')
if labels.ndim == 1:
labels = np.reshape(labels, (1, -1))
labels[:, 0] = labels[:, 0] + self.keden_labels_num
labels = pad_targets(labels, pad_dim)
else:
# image_tensor = torch.zeros(3, resize_dim, resize_dim)
labels = np.zeros((pad_dim, 5))
return image_tensor, labels
wematan
harpreetsahota204
BloodAxe
legenda971
💡 Your Question
Hi,
I'm trying to train YOLO NAS on a custom dataset, that basically consists of two different datasets with 17 classes in the first one and 12 new classes in the second one. My task is to take the whole first dataset and add 10% of the second dataset to it (since the second one contains new classes but has far more images so I don't want to take all of them, but just a random portion) and treat it as a single input dataset for the model. I wrote two custom Dataset classes and actually already have posted a question #1056 about varying number of bboxes, and your suggestion on padding helped a lot. Thanks!
And now I have my datasets (I sampled 10% of the second dataset using torch.utils.data.Subset()), which I later combine in one dataset using torch.utils.data.ConcatDataset(), but there is another unexpected issue: the model starts training just okay but after several epochs I see that practically all losses and metrics drop to zero.
Also, I am using a pre-trained COCO weights. I've seen losses go to NaN when the learning rate is too high, but I've never seen all the metrics to be zero. I tried to decrease the learning rate (my initial lr was 5e-4, I tried 5e-5, 5e-6), removed warmup epochs, and I also tried to train the model from scratch without using COCO weights - the model eventually goes to zeros in any of those cases.
Now I start to suspect that I did something wrong in either Dataset() or Dataloader(), but I just can't understand what is wrong. So your help or any useful debugging suggestions will be much appreciated.
Versions
No response