def test_training_step(self):
# Setup the DataModule
data_path = 'data-bin'
train_dataset = data_helper.get_dataset(data_root=data_path, mode='train')
val_dataset = data_helper.get_dataset(data_root=data_path, mode='val')
data_module = DetectionDataModule(train_dataset, val_dataset, batch_size=16)
# Load model
model = yolov5s()
model.train()
# Trainer
trainer = pl.Trainer(max_epochs=1)
trainer.fit(model, data_module)
def test_training_step():
# Setup the DataModule
data_path = "data-bin"
train_dataset = data_helper.get_dataset(data_root=data_path, mode="train")
val_dataset = data_helper.get_dataset(data_root=data_path, mode="val")
data_module = DetectionDataModule(train_dataset, val_dataset, batch_size=8)
# Load model
model = DefaultTask(arch="yolov5n")
model = model.train()
# Trainer
trainer = pl.Trainer(max_epochs=1)
trainer.fit(model, data_module)
def test_get_dataset(self):
# Acquire the images and labels from the coco128 dataset
train_dataset = data_helper.get_dataset(data_root='data-bin', mode='train')
# Test the datasets
image, target = next(iter(train_dataset))
self.assertIsInstance(image, Tensor)
self.assertIsInstance(target, Dict)
def test_get_dataset():
# Acquire the images and labels from the coco128 dataset
train_dataset = data_helper.get_dataset(data_root="data-bin", mode="train")
# Test the datasets
image, target = next(iter(train_dataset))
assert isinstance(image, Tensor)
assert isinstance(target, dict)
def test_detection_data_module(self):
# Setup the DataModule
batch_size = 4
train_dataset = data_helper.get_dataset(data_root='data-bin', mode='train')
data_module = DetectionDataModule(train_dataset, batch_size=batch_size)
self.assertEqual(data_module.batch_size, batch_size)
data_loader = data_module.train_dataloader(batch_size=batch_size)
images, targets = next(iter(data_loader))
self.assertEqual(len(images), batch_size)
self.assertIsInstance(images[0], Tensor)
self.assertEqual(len(images[0]), 3)
self.assertEqual(len(targets), batch_size)
self.assertIsInstance(targets[0], Dict)
self.assertIsInstance(targets[0]["image_id"], Tensor)
self.assertIsInstance(targets[0]["boxes"], Tensor)
self.assertIsInstance(targets[0]["labels"], Tensor)
def test_detection_data_module():
# Setup the DataModule
batch_size = 4
train_dataset = data_helper.get_dataset(data_root="data-bin", mode="train")
data_module = DetectionDataModule(train_dataset, batch_size=batch_size)
assert data_module.batch_size == batch_size
data_loader = data_module.train_dataloader()
images, targets = next(iter(data_loader))
assert len(images) == batch_size
assert isinstance(images[0], Tensor)
assert len(images[0]) == 3
assert len(targets) == batch_size
assert isinstance(targets[0], dict)
assert isinstance(targets[0]["image_id"], Tensor)
assert isinstance(targets[0]["boxes"], Tensor)
assert isinstance(targets[0]["labels"], Tensor)