(self)
| 16 | |
| 17 | class TestPytorchGraph(unittest.TestCase): |
| 18 | def test_graph(self): |
| 19 | model = torchvision.models.vgg16() |
| 20 | g = hl.build_graph(model, torch.zeros([1, 3, 224, 224])) |
| 21 | g.save(os.path.join(OUTPUT_DIR, "pytorch_vgg16.pdf")) |
| 22 | |
| 23 | model = torchvision.models.resnet50() |
| 24 | g = hl.build_graph(model, torch.zeros([1, 3, 224, 224])) |
| 25 | g.save(os.path.join(OUTPUT_DIR, "pytorch_resnet50.pdf")) |
| 26 | |
| 27 | # Clean up |
| 28 | shutil.rmtree(OUTPUT_DIR) |
| 29 | |
| 30 | |
| 31 | def test_resnet_blocks(self): |