def test_train_no_epoch(self):
self.assertRaises(
TypeError,
lambda: tiny_yolo2_train.train(model=self.model,
train_data=self.dataset,
epochs=None,
batch_size=self.batch_size))
def test_train(self):
ckpt, pb = tiny_yolo2_train.train(model=self.model,
train_data=self.dataset,
epochs=self.epochs,
batch_size=self.batch_size)
file_list = ["tiny_yolo2.ckpt", "tiny_yolo2.pb"]
self.assertIn(file_list[0], ckpt)
self.assertIn(file_list[1], pb)
self.assertTrue(
tests.is_dir_contains(dirs=self.model_path, file_list=file_list))
def test_train_no_model(self):
self.assertRaises(
AttributeError, lambda: tiny_yolo2_train.train(
model=None, train_data=self.dataset, epochs=None, batch_size=1)
)