|
2 | 2 | import os
|
3 | 3 |
|
4 | 4 | # Use system installed Python packages
|
5 |
| -PYT_PATH='/opt/conda/lib/python3.8/site-packages' |
| 5 | +PYT_PATH='/opt/conda/lib/python3.8/site-packages' if not 'PYT_PATH' in os.environ else os.environ["PYT_PATH"] |
6 | 6 |
|
7 | 7 | # Root directory for torch_tensorrt. Set according to docker container by default
|
8 |
| -TOP_DIR='/opt/pytorch/torch_tensorrt' |
| 8 | +TOP_DIR='/torchtrt' if not 'TOP_DIR' in os.environ else os.environ["TOP_DIR"] |
9 | 9 |
|
10 | 10 | # Download the dataset
|
11 | 11 | @nox.session(python=["3"], reuse_venv=True)
|
12 | 12 | def download_datasets(session):
|
| 13 | + print("Downloading dataset to path", os.path.join(TOP_DIR, 'examples/int8/training/vgg16')) |
13 | 14 | session.chdir(os.path.join(TOP_DIR, 'examples/int8/training/vgg16'))
|
14 |
| - session.run_always('wget', 'https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz') |
15 |
| - session.run_always('tar', '-xvzf', 'cifar-10-binary.tar.gz') |
| 15 | + session.run_always('wget', 'https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz', external=True) |
| 16 | + session.run_always('tar', '-xvzf', 'cifar-10-binary.tar.gz', external=True) |
16 | 17 | session.run_always('mkdir', '-p',
|
17 |
| - os.path.join(TOP_DIR, 'tests/accuracy/datasets/data')) |
| 18 | + os.path.join(TOP_DIR, 'tests/accuracy/datasets/data'), |
| 19 | + external=True) |
18 | 20 | session.run_always('cp', '-rpf',
|
19 | 21 | os.path.join(TOP_DIR, 'examples/int8/training/vgg16/cifar-10-batches-bin'),
|
20 | 22 | os.path.join(TOP_DIR, 'tests/accuracy/datasets/data/cidar-10-batches-bin'),
|
|
0 commit comments