def test_griewank():
for graph_length in [2**i for i in range(1, 6)]:
g = gen_linear_graph(graph_length)
assert g.size_fwd == graph_length
total_cost = sum(g.cost_ram.values())
scheduler_result = solve_griewank(g, total_cost)
assert scheduler_result.feasible
import pandas as pd
import matplotlib.pyplot as plt
if __name__ == "__main__":
N = 16
for B in range(4, 12):
# model = get_keras_model("MobileNet")
# g = dfgraph_from_keras(mod=model)
g = gen_linear_graph(N)
scratch_dir = remat_data_dir() / f"scratch_linear" / str(N) / str(B)
scratch_dir.mkdir(parents=True, exist_ok=True)
data = []
scheduler_result_all = solve_checkpoint_all(g)
scheduler_result_sqrtn = solve_chen_sqrtn(g, True)
scheduler_result_griewank = solve_griewank(g, B)
plot(scheduler_result_all,
False,
save_file=scratch_dir / "CHECKPOINT_ALL.png")
plot(scheduler_result_sqrtn,
False,
save_file=scratch_dir / "CHEN_SQRTN.png")
plot(scheduler_result_griewank,
False,
save_file=scratch_dir / "GRIEWANK.png")
data.append({
"Strategy":
str(scheduler_result_all.solve_strategy.value),
"Name":
"CHECKPOINT_ALL",
"CPU":
N = args.num_layers
IMPOSED_SCHEDULE = args.imposed_schedule
APPROX = False
EPS_NOISE = 0
SOLVE_R = False
# Compute integrality gap for each budget
for B in reversed(range(4, N + 3)): # Try several budgets
g = gen_linear_graph(N)
scratch_dir = remat_data_dir(
) / f"scratch_integrality_gap_linear" / f"{N}_layers" / str(
IMPOSED_SCHEDULE) / f"{B}_budget"
scratch_dir.mkdir(parents=True, exist_ok=True)
data = []
griewank = solve_griewank(g, B)
logging.info("--- Solving LP relaxation for lower bound")
lb_lp = lower_bound_lp_relaxation(g,
B,
approx=APPROX,
eps_noise=EPS_NOISE,
imposed_schedule=IMPOSED_SCHEDULE)
plot(lb_lp, False, save_file=scratch_dir / "CHECKMATE_LB_LP.png")
logging.info("--- Solving ILP")
ilp = solve_ilp_gurobi(g,
B,
approx=APPROX,
eps_noise=EPS_NOISE,
imposed_schedule=IMPOSED_SCHEDULE,
if model_name not in CHAIN_GRAPH_MODELS:
logger.info(
f"Running Chen's greedy baseline (no AP) as model is non-linear")
futures = [
remote_solve_chen_greedy(g, float(b), True)
for b in greedy_eval_points
]
result_dict[SolveStrategy.CHEN_SQRTN_NOAP] = get_futures(
list(futures), desc="Greedy (No AP)")
# sweep griewank baselines
if model_name in CHAIN_GRAPH_MODELS:
logger.info(f"Running Griewank baseline (APs only)")
clean_griewank_cache()
solve_griewank(
g, 1
) # prefetch griewank solution from s3, otherwise ray will cause race condition
griewank_eval_points = range(1, g.size + 1)
remote_solve_griewank = ray.remote(num_cpus=1)(solve_griewank).remote
futures = [
remote_solve_griewank(g, float(b)) for b in griewank_eval_points
]
result_dict[SolveStrategy.GRIEWANK_LOGN] = get_futures(
list(futures), desc="Griewank (APs only)")
# sweep optimal ilp baseline
if not args.skip_ilp:
ilp_log_base = log_base / "ilp_log"
ilp_log_base.mkdir(parents=True, exist_ok=True)
# todo load any ILP results from cache
remote_ilp = ray.remote(