def test_label_has_ship():
dataset = Ships(15)
for i in range(len(dataset)):
sample = dataset[i]["target"]
if sample[0] == 1:
assert ~torch.isnan(sample[1]).item()
elif sample[0] == 0:
assert torch.isnan(sample[1]).item()
else:
assert False
def train_dataloader(opt):
print("[INFO]: Train dataloader called")
dataset = Ships(n_samples=opt.train_len)
sampler = None
shuffle = True
loader = DataLoader(dataset=dataset,
batch_size=opt.batch_size,
num_workers=opt.num_workers,
pin_memory=opt.pin_memory,
sampler=sampler,
shuffle=shuffle)
print("samples - ", len(dataset))
return loader
def test_dataloader(opt: Namespace) -> DataLoader:
print("[INFO]: Test dataloader called")
dataset = Ships(n_samples=opt.test_len)
sampler = None
shuffle = True
loader = DataLoader(
dataset=dataset,
batch_size=opt.batch_size,
num_workers=opt.num_workers,
pin_memory=opt.pin_memory,
sampler=sampler,
shuffle=shuffle,
)
print("samples - ", len(dataset))
return loader