Sign inSign up

jaehun/ost

By jaehun

•Updated over 4 years ago

One-Shot Tuner for Deep Learning Compilers

Image
0

134

jaehun/ost repository overview

⁠prior-guided task sampilng(PBS) and Exploration Based code Sampling(EBS)

The extracted dataset for CUDA data is included.

  • -p activates Prior Guided Task Sampling.
  • -e activates Exploration Based code sampling.
python3 dataset_generate/sampling.py -p -e

⁠Training a cost model

python3 train_model/train.py --dataset_dir <dataset path (e.g. /root/ost/dataset_generate)> --layout NCHW --batch 1

⁠Evaluating and collecting results

You can run all main experiment(CUDA device,NCHW format batch 1) using main.sh script.

main.sh

Create a folder using the save path and parameters that can be specified in the script and save the result (second,flops/s and end-to-end time).

Example path

  • /root/eval_tuner/save_path/resnet-18/NCHW/1/sa/flops.npy

⁠Docker guide

If setting the environment is difficult, try using Docker container

docker run -it --gpus 1 --rm  jaehun/ost:v2 bash # docker running
cd /root/tvm
./main.sh 									  # start experiment
python3 get_result.py                         # get results

Tag summary

Content type

Image

Digest

Size

8.1 GB

Last updated

over 4 years ago

docker pull jaehun/ost:v2