from experiments.common.graph_plotting import plot
from remat.core.solvers.strategy_checkpoint_all import solve_checkpoint_all
from remat.core.solvers.strategy_chen import solve_chen_sqrtn
from remat.core.solvers.strategy_griewank import solve_griewank
from remat.core.solvers.strategy_optimal_ilp import solve_ilp_gurobi
from remat.core.utils.timer import Timer
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")
from experiments.common.definitions import remat_data_dir
import numpy as np
import pandas as pd
import glob
import re
# compute aggregated tables of max and geomean lp approximation ratios
exp_name_re = re.compile(r"^(?P<platform>.+?)_(?P<model_name>.+?)_(?P<batch_size>[0-9]+?)_(?P<input_shape>None|.+?)$")
dfs = []
for path in (remat_data_dir() / 'budget_sweep').glob('**/slowdowns.csv'):
slowdown_df = pd.read_csv(path)
matches = exp_name_re.match(path.parents[0].name)
model_name = matches.group('model_name')
slowdown_df['Model name'] = [model_name] * len(slowdown_df)
dfs.append(slowdown_df)
df = pd.concat(dfs)
del df['Unnamed: 0']
for valuekey in ['geomean_slowdown', 'max']:
pivot_df = pd.pivot_table(df, values=valuekey, index=['Model name'], columns=['method'])
pivot_df.to_csv(remat_data_dir() / 'budget_sweep' / f"{valuekey}_aggr.csv")
# compute lp relaxation speedups
ilp_runtime_dict = {}
lp_runtime_dict = {}
for model in ['p32xlarge_vgg_unet_32_None', 'p32xlarge_ResNet50_256_None', 'p32xlarge_MobileNet_512_None', 'p32xlarge_VGG16_256_None', 'p32xlarge_VGG19_256_None']:
ilp_matcher = re.compile(r"Explored [0-9]+ nodes \([0-9]+ simplex iterations\) in (?P<ilp_runtime>[0-9\.]+) seconds")
lp_matcher = re.compile(r"Solved in [0-9]+ iterations and (?P<lp_runtime>[0-9\.]+) seconds")
ilp_runtimes = []
for path in (remat_data_dir() / 'budget_sweep' / model / 'ilp_log').glob('./*.log'):
with path.open('r') as f:
return results
if __name__ == "__main__":
logger = logging.getLogger("budget_sweep")
logger.setLevel(logging.DEBUG)
# due to bug on havoc, limit parallelism on high-core machines
if os.cpu_count() > 48:
os.environ["OMP_NUM_THREADS"] = "1"
args = extract_params()
ray.init(temp_dir="/tmp/ray_checkpoint", redis_password=str(uuid.uuid1()), num_cpus=os.cpu_count(),
object_store_memory=1024 * 1024 * 1024 if os.cpu_count() < 48 else None) # include_webui=args.debug
key = "_".join(map(str, [args.platform, args.model_name, args.batch_size, args.input_shape]))
log_base = remat_data_dir() / "budget_sweep" / key
####
# Begin budget_sweep data collection
####
model_name = args.model_name
# load costs, and plot optionally, if platform is not flops
logger.info(f"Loading costs")
if args.platform == "flops":
cost_model = None
else:
cost_model = CostModel(model_name, args.platform, log_base, quantization=5)
cost_model.fit()
# gen redis key
if __name__ == "__main__":
logging.basicConfig(level=logging.DEBUG)
# Set parameters
args = parse_args()
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")
_args = parser.parse_args()
_args.input_shape = _args.input_shape if _args.input_shape else None
return _args
if __name__ == "__main__":
logging.getLogger().setLevel(logging.INFO)
# due to bug on havoc, limit parallelism on high-core machines
if os.cpu_count() > 48:
os.environ["OMP_NUM_THREADS"] = "1"
args = extract_params()
key = "_".join(map(str,
[args.platform, args.model_name, args.input_shape]))
log_base = remat_data_dir() / "max_batch_size" / key
shutil.rmtree(log_base, ignore_errors=True)
pathlib.Path(log_base).mkdir(parents=True, exist_ok=True)
result_dict: Dict[int, Dict[SolveStrategy,
List[ScheduledResult]]] = defaultdict(
lambda: defaultdict(list))
model_name = args.model_name
# load costs, and plot optionally, if platform is not flops
logging.info(f"Loading costs")
if args.platform == "flops":
cost_model = None
else:
cost_model = CostModel(model_name,
args.platform,
log_base,