Method run_with

caffe2/python/checkpoint_test.py:68–95

(self, builder)

class TestCheckpoint(TestCase):
    def run_with(self, builder):
        with Cluster():
            with Job() as job:
                outputs = build_pipeline(node_id=0)
            output_fetcher = Task(step=core.Net('empty'), outputs=outputs)

            def fetch_total(session):
                session.run(output_fetcher)
                return output_fetcher.outputs()[0].fetch()

            session, checkpoint = builder()
            job.compile(LocalSession)
            num_epochs = JobRunner(job, checkpoint).train(session)
            self.assertEqual(num_epochs, len(EXPECTED_TOTALS))
            self.assertEqual(fetch_total(session), EXPECTED_TOTALS[-1])

            for initial_epoch in range(1, num_epochs + 1):
                session, checkpoint = builder()
                JobRunner(
                    job,
                    checkpoint, resume_from_epoch=initial_epoch
                ).train(session)
                self.assertEqual(fetch_total(session), EXPECTED_TOTALS[-1])

            for epoch in range(1, num_epochs + 1):
                session.run(checkpoint.load(epoch))
                self.assertEqual(fetch_total(session),
                                  EXPECTED_TOTALS[epoch - 1])

    def test_single_checkpoint(self):
        # test single node