Hi all!
I have implemented some random triplet sampler in the following way:
class TripletSampler(Sampler):
def __init__(self, dataset, batch_size=1, drop_last=True):
self.dataset = dataset
self.batch_size = batch_size
self.shuffle = shuffle
self.drop_last = drop_last
self.triplets = None
self.sampler = SequentialSampler(self.dataset)
def __iter__(self):
labels = self.dataset.labels
batch = []
triplets = []
for r_idx in self.sampler:
p_idx = np.random.choice(np.where(labels == labels[r_idx])[0], 1)[0]
n_idx = np.random.choice(np.where(labels != labels[r_idx])[0], 1)[0]
batch.append([r_idx, p_idx, n_idx])
if len(batch) == self.batch_size:
yield batch
triplets.append(batch)
batch = []
gc.collect()
if len(batch) > 0 and not self.drop_last:
yield batch
triplets.append(batch)
self.triplets = triplets
def __len__(self):
if self.drop_last:
return len(self.sampler) // self.batch_size
else:
return (len(self.sampler) + self.batch_size - 1) // self.batch_size
While I am training the model, the GPU memory usage slowly increases until it has not enough memory to continue running. If I do not use the sampler and instead pre-calculate the triplets and then feed it to the dataset it works fine.
I tried adding some gc.colllector() but it memory is still increasing at a really slow rate. Any idea on why is this happening?
If needed, this is the implementation of the __getitem__ method inside the dataset (hope it is self-understandable):
def __getitem__(self, index):
# get the triplet indexes and use them to obtain their paths
r_idx, p_idx, n_idx = index
R_path, P_path, N_path = self.data_path[r_idx], self.data_path[p_idx], self.data_path[n_idx]
# load the images in the triplet
R, P, N = data_utils.read_exr(R_path), data_utils.read_exr(P_path), data_utils.read_exr(N_path)
# load the labels of the images with its names
R_l, P_l, N_l = self.labels[r_idx], self.labels[p_idx], self.labels[n_idx]
# apply transforms to all the images in the triplet
if self.transforms is not None:
R, P, N = self.transforms(R), self.transforms(P), self.transforms(N)
# return the images in the triplet
return R, P, N, R_l, P_l, N_l