Skip to content

Instantly share code, notes, and snippets.

@tangzhankun
Last active January 31, 2019 10:29
Show Gist options
  • Save tangzhankun/c7e87b43a62ae442e03cb21005e51272 to your computer and use it in GitHub Desktop.
Save tangzhankun/c7e87b43a62ae442e03cb21005e51272 to your computer and use it in GitHub Desktop.
#!/usr/bin/python
import itertools
import subprocess
import csv
with open('tensorflow_bench_result.csv', 'wb') as csvFile:
csvWriter = csv.writer(csvFile, quoting=csv.QUOTE_ALL)
#gpus = ['0', '1', '2', '3', '4', '5', '6', '7']
#batch_size = ['32', '64', '128']
#model = ['alexnet', 'resnet50', 'vgg16', 'inception3']
gpus = ['0', '1', '2', '3']
batch_size = ['32']
model = ['alexnet', 'resnet50']
for L in range(2, len(gpus)):
for subset in set(itertools.combinations(gpus, L)):
chosenGPUs = ','.join(subset)
print(chosenGPUs)
for i in range(len(batch_size)):
chosenBatchSize = batch_size[i]
print('\t' + chosenBatchSize)
for j in range(len(model)):
chosenModel = model[j]
print('\t\t' + chosenModel)
output = subprocess.check_output(['docker', 'run', '--rm', '-it', '--runtime=nvidia', '-e', 'NVIDIA_VISIBLE_DEVICES=' + chosenGPUs, 'tangzhankun/tensorflow_benchmark:topo', '/notebooks/run_benchmark.sh', str(L), chosenBatchSize, chosenModel])
lines = output.split('\n')
FPSline = lines[len(lines)-3];
FPStoken = FPSline.split(':')
FPS = FPStoken[1].strip()
csvLine = []
csvLine.append(chosenModel)
csvLine.append(chosenBatchSize)
csvLine.append('-'.join(subset))
csvLine.append(unicode(FPS, "utf-8"))
#csvWriter.writerow(['Spam', 'Lovely Spam', 'Wonderful Spam'])
csvWriter.writerow(csvLine)
csvFile.flush()
print('\t\t\t csv line:' + ",".join(csvLine))
#docker run --rm -it --runtime=nvidia -e "NVIDIA_VISIBLE_DEVICES=0,1,2,3" tangzhankun/tensorflow_benchmark:topo /notebooks/run_benchmark.sh 2 32 resnet50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment