def cli_main():
pl.seed_everything(1234)
# ------------
# args
# ------------
parser = ArgumentParser()
parser = pl.Trainer.add_argparse_args(parser)
parser.add_argument("--ray_accelerator_num_workers", type=int, default=4)
parser.add_argument("--ray_accelerator_cpus_per_worker",
type=int,
default=1)
parser.add_argument("--use_gpu", type=bool, default=False)
parser = LitMNIST.add_model_specific_args(parser)
parser = MNISTDataModule.add_argparse_args(parser)
args = parser.parse_args()
ray.init(address='auto')
print('''This cluster consists of
{} nodes in total
{} CPU resources in total
'''.format(len(ray.nodes()),
ray.cluster_resources()['CPU']))
# ------------
# data
# ------------
dm = MNISTDataModule(data_dir='',
val_split=5000,
num_workers=1,
normalize=False,
seed=42,
batch_size=32)
# ------------
# model
# ------------
model = LitMNIST(args.hidden_dim, args.learning_rate)
# ------------
# training
# ------------
# ray.init()
plugin = RayPlugin(
num_workers=args.ray_accelerator_num_workers,
num_cpus_per_worker=args.ray_accelerator_cpus_per_worker,
use_gpu=args.use_gpu)
trainer = pl.Trainer(gpus=args.gpus,
precision=args.precision,
plugins=[plugin],
max_epochs=args.max_epochs)
trainer.fit(model, datamodule=dm)
# ------------
# testing
# ------------
result = trainer.test(model, datamodule=dm)
pprint(result)
def test_anomaly(self):
anomaly = 9
datamodule = MNISTDataModule(self.DATA_ROOT, exclude=anomaly)
datamodule.prepare_data()
datamodule.setup()
with self.subTest(split='train'):
self._check_excluded(anomaly, datamodule.train_dataloader())
with self.subTest(split='val'):
self._check_included(anomaly, datamodule.val_dataloader())
with self.subTest(split='test'):
self._check_included(anomaly, datamodule.test_dataloader())
def test_accuracy_returned_on_test(self):
datamodule = MNISTDataModule(data_dir='../data')
trainer = pl.Trainer(logger=False)
test_results, *_ = trainer.test(self.net, datamodule=datamodule)
self.assertIsNotNone(test_results)
self.assertLessEqual(0, test_results['test/accuracy'])
self.assertGreaterEqual(1, test_results['test/accuracy'])
def test_train_size(self):
train_size = 500
datamodule = MNISTDataModule(self.DATA_ROOT, train_size=train_size)
datamodule.prepare_data()
datamodule.setup()
train_data = datamodule.train_dataloader().dataset
self.assertEqual(train_size, len(train_data))
class TestLatent(unittest.TestCase):
def setUp(self):
self.data = MNISTDataModule(data_dir='../data')
self.data.prepare_data()
self.data.setup()
self.net = build_ae('vae', self.data.dims)
self.latent_viz = Latent(self.net)
def test_sample(self):
samples = self.latent_viz.sample(16)
expected_shape = torch.Size((3, 2 * 32 + 3 * 2, 8 * 32 + 9 * 2)) # 2 rows x 8 cols + 2 pad each side
self.assertEqual(expected_shape, samples.shape)
def test_reconstruct(self):
_, comparison = self.latent_viz.reconstruct(self.data, num_comparison=32)
expected_shape = torch.Size((3, 32 * 32 + 33 * 2, 2 * 32 + 3 * 2)) # 32 rows x 2 cols + 2 pad each side
self.assertEqual(expected_shape, comparison.shape)
def test_interpolate(self):
data_iter = iter(self.data.test_dataloader())
start_batch, _ = next(data_iter)
end_batch, _ = next(data_iter)
steps = 1
interpolations = self.latent_viz.interpolate(start_batch, end_batch, steps=steps)
expected_shape = (steps + 2, 3, 4 * 32 + 5 * 2, 8 * 32 + 9 * 2) # 3 steps x 4 rows x 8 cols + 2 pad each side
self.assertEqual(expected_shape, interpolations.shape)
def test_reduce(self):
reduced_latents, labels = self.latent_viz.reduce(self.data.val_dataloader())
self.assertEqual(len(self.data.mnist_val), reduced_latents.shape[0])
self.assertEqual(len(self.data.mnist_val), labels.shape[0])
class TestAnomalyDetection(unittest.TestCase):
def setUp(self):
self.data = MNISTDataModule(data_dir='../data', exclude=9)
self.data.prepare_data()
self.data.setup()
self.net = build_ae('vanilla', self.data.dims)
self.anomaly_detector = AnomalyDetection(self.net)
def test_anomaly_labels(self):
class_labels = torch.randint(0, 10, size=(100,))
dummy_features = torch.randn(100, 50)
dummy_dataset = TensorDataset(dummy_features, class_labels)
expected_anomaly_labels = (class_labels == 9).numpy().astype(np.int)
actual_anomaly_labels = self.anomaly_detector.get_test_anomaly_labels(DataLoader(dummy_dataset, batch_size=32),
anomaly_value=9)
self.assertListEqual(expected_anomaly_labels.tolist(), actual_anomaly_labels.tolist())
def test_score(self):
self.anomaly_detector.autoencoder.forward = mock.MagicMock(wraps=lambda x: torch.zeros_like(x))
test_dataloader = self.data.test_dataloader()
scores = self.anomaly_detector.score(test_dataloader)
self.assertIsInstance(scores, np.ndarray)
self.assertEqual(10000, scores.shape[0]) # number of scores is length of test data
self.assertTrue(np.all(scores >= 0.), msg="Regression: scores have to be all non-negative.")
def test_roc(self):
tpr, fpr, thresholds, auc = self.anomaly_detector.get_test_roc(self.data)
test_dataloader = self.data.test_dataloader()
scores = self.anomaly_detector.score(test_dataloader)
self.assertTrue(np.all(thresholds >= np.min(scores)))
self.assertTrue(np.all(thresholds[1:] <= np.max(scores)))
self.assertGreaterEqual(thresholds[1], np.max(scores))
self.assertLessEqual(0, auc)
self.assertGreaterEqual(1, auc)
def training_step(self, batch, *args, **kwargs):
x, y = batch
logits = self.forward(x)
loss = self._loss_fn(logits, y)
self.log('train_loss', loss)
return {'loss': loss}
def validation_step(self, batch, batch_idx, *args, **kwargs):
x, y = batch
logits = self.forward(x)
loss = self._loss_fn(logits, y)
self.log('val_loss', loss)
pred = logits.argmax(dim=1, keepdim=True)
accuracy = pred.eq(y.view_as(pred)).sum().item() / len(x)
return {'val_loss': loss, 'accuracy': accuracy}
def configure_optimizers(self):
optimizer = torch.optim.Adam(self.parameters(),
lr=self.hparams.lr)
return optimizer
if __name__ == '__main__':
model = LightningMNISTClassifier(lr=1e-3)
trainer = pl.Trainer(max_epochs=2, default_root_dir='/tmp/lightning')
dm = MNISTDataModule('https://s3-us-west-2.amazonaws.com/determined-ai-test-data/pytorch_mnist.tar.gz')
trainer.fit(model, datamodule=dm)
def _get_mnist(self,
data_root,
batch_size=32,
exclude=None,
train_size=None):
return MNISTDataModule(data_root, batch_size, train_size, exclude)
def setUp(self):
self.data = MNISTDataModule(data_dir='../data')
self.data.prepare_data()
self.data.setup()
self.net = build_ae('vae', self.data.dims)
self.latent_viz = Latent(self.net)
def setUp(self):
self.data = MNISTDataModule(data_dir='../data', exclude=9)
self.data.prepare_data()
self.data.setup()
self.net = build_ae('vanilla', self.data.dims)
self.anomaly_detector = AnomalyDetection(self.net)