def plot_runtime_comparison(base, heuristic, bound, out_file):
"""Compare different runtime settings on the same budget and heuristic."""
base_dir = util.get_output_path(ASYMPTOTICS_MOD + '/' + base, '')
results = []
paths = glob.glob(base_dir + '*.json')
for result_path in paths:
js = None
with open(result_path, 'r') as jf:
js = json.loads(jf.read())
assert js is not None
# if js['config']['runtime'] != 'V2': continue
if js['config']['heuristic'] != str(heuristic) and \
js['config']['heuristic'] != str(heuristic) + ' (Unoptimized)':
print('skipping "{}", since {} != {}'.format(
result_path, js['config']['heuristic'], str(heuristic)))
continue
if js['config']['memory'] != str(bound):
print('skipping "{}", since {} != {}'.format(
result_path, js['config']['memory'], str(bound)))
continue
assert len(results) == 0 or results[-1]['config']['releases'] == js[
'config']['releases']
assert len(results) == 0 or results[-1]['config']['kwargs'] == js[
'config']['kwargs']
results.append(js)
for res in results:
v1_banishing = res['config']['runtime'] == 'V1'
if not v1_banishing:
eager_evict = 'eager_evict' in res['config']['runtime_features']
if eager_evict:
runtime_label = 'Eager eviction'
else:
runtime_label = 'No eager eviction'
else:
runtime_label = 'Banishing'
l = plt.plot(res['results']['layers'],
res['results']['computes'],
label=runtime_label,
alpha=0.7,
marker='X')
l_color = l[0].get_color()
if res['results']['had_OOM']:
plt.axvline(x=res['results']['layers'][0], color=l_color, linestyle='dotted', \
alpha=0.3, linewidth=3)
if res['results']['had_thrash']:
plt.axvline(x=res['results']['layers'][-1], color=l_color, linestyle='--', \
alpha=0.3, linewidth=3)
if isinstance(bound, Log2Bound):
plot_treeverse(results[0]['config']['ns'])
else:
plot_tq(results[0]['config']['ns'])
plt.grid(True)
plt.xlabel(r'Number of Layers $n$')
plt.ylabel(r'Additional Compute')
plt.legend()
plt.title('Layers vs. Compute Overhead\n{} Heuristic, {} Memory'.format(
str(heuristic), str(bound)))
plt.savefig(base_dir + out_file, dpi=300)
plt.clf()
def run_asymptotics(base,
ns,
heuristic,
bound,
runtime,
releases=True,
**kwargs):
config = {
'ns': ns,
'heuristic': str(heuristic),
'heuristic_features': list(heuristic.FEATURES),
'memory': str(bound),
'releases': releases,
'runtime': runtime.ID,
'runtime_features': list(runtime.FEATURES),
'kwargs': kwargs
}
p = Pool()
print('generating asymptotics data for config: {}...'.format(
json.dumps(config, indent=2)))
args = []
for n in ns:
args.append([n, bound(n), heuristic, runtime, releases, kwargs])
t = time.time()
rts = p.map(run, *zip(*args))
t = time.time() - t
succ_ns, succ_rts = chop_failures(ns, rts)
print(' - succeeded between n={} and n={}'.format(succ_ns[0],
succ_ns[-1]))
print(' done, took {} seconds.'.format(t))
results = {
'layers':
succ_ns,
'computes':
list(map(lambda rt: rt.telemetry.summary['remat_compute'], rts)),
'had_OOM':
ns[0] != succ_ns[0],
'had_thrash':
ns[-1] != succ_ns[-1]
}
date_str = datetime.now().strftime('%Y%m%d-%H%M%S-%f')
base_mod = ASYMPTOTICS_MOD + '/' + base
out_file = '{}-{}-{}.json'.format(date_str, heuristic.ID, bound.ID)
util.ensure_output_path(base_mod)
out_path = util.get_output_path(base_mod, out_file)
with open(out_path, 'w') as out_f:
out_f.write(
json.dumps({
'config': config,
'results': results
}, indent=2))
print('-> done, saved to "{}"'.format(out_path))
def plot_tq_local(n=80):
bound = TQBound()
heuristic = DTRLocal()
runtime = RuntimeV2EagerOptimized
rt_kwargs = {'stats': True, 'trace': True}
rt = run(n,
bound(n),
heuristic,
runtime,
releases=True,
rt_kwargs=rt_kwargs)
render_trace(n, rt.telemetry, dedup_grad=True)
plt.legend()
plt.title('$n = {}$, $B$ = {}, DTR-Local Heuristic'.format(n, str(bound)))
plt.savefig(util.get_output_path(TRACE_MOD, 'tq_local.png'), dpi=300)
plt.clf()
def plot_heuristic_comparison(base, bound, runtime, out_file):
"""Compare different heuristics on the same runtime and budget."""
base_dir = util.get_output_path(ASYMPTOTICS_MOD + '/' + base, '')
results = []
for result_path in glob.glob(base_dir + '*.json'):
js = None
with open(result_path, 'r') as jf:
js = json.loads(jf.read())
assert js is not None
if js['config']['runtime'] != runtime.ID: continue
if js['config']['memory'] != str(bound): continue
assert len(results) == 0 or results[-1]['config']['releases'] == js[
'config']['releases']
assert len(
results) == 0 or results[-1]['config']['ns'] == js['config']['ns']
results.append(js)
for res in results:
h_str = res['config']['heuristic']
h_color, h_marker = HEURISTICS[h_str].COLOR, HEURISTICS[h_str].MARKER
plt.plot(res['results']['layers'], res['results']['computes'], label=h_str, \
color=h_color, marker=h_marker, alpha=0.5, linewidth=3, ms=10)
if res['results']['had_OOM']:
plt.axvline(x=res['results']['layers'][0], color=h_color, linestyle='dotted', \
alpha=0.3, linewidth=3)
if res['results']['had_thrash'] is not None:
plt.axvline(x=res['results']['layers'][-1], color=h_color, linestyle='--', \
alpha=0.3, linewidth=3)
if isinstance(bound, Log2Bound):
plot_treeverse(results[0]['config']['ns'])
else:
plot_tq(results[0]['config']['ns'])
plt.grid(True)
plt.xlabel(r'Number of Layers $n$')
plt.ylabel(r'Additional Compute')
plt.legend()
plt.title('Layers vs. Compute Overhead ({} Memory)'.format(str(bound)))
plt.savefig(base_dir + out_file, dpi=300)
plt.clf()