def test_checkmate_to_simrd_analytical_cost():
if not have_checkmate:
return
test_log_filename = 'data/checkmate_simrd.log'
batch_size = 1
with open(test_log_filename, 'w') as test_log:
for name in MODEL_NAMES:
try:
model = get_keras_model(name)
dfg = dfgraph_from_keras(model,
batch_size=batch_size,
loss_cpu_cost=0,
loss_ram_cost=(4 * batch_size))
g = from_dfgraph(dfg)
rt, result, pr = run_baseline(g.get_closure(),
stats=True,
trace=True)
print('Baseline simrd results for {}:'.format(name),
file=test_log)
print(json.dumps(result, indent=2), file=test_log)
print(file=test_log)
except Exception as e:
print('Failed for {}:'.format(name), file=test_log)
print(traceback.format_exc(), file=test_log)
print(file=test_log)
print(
'saved Checkmate -> simrd test log to [{}]'.format(test_log_filename))
def test_testnet_checkpointall():
model = get_keras_model("test")
g = dfgraph_from_keras(mod=model)
assert g.size_fwd == 6
scheduler_result = solve_checkpoint_all(g)
assert scheduler_result.feasible
assert scheduler_result.schedule_aux_data.cpu == sum(g.cost_cpu.values())
def test_testnet_optimalilp():
try:
import gurobipy as _
except ImportError as e:
logging.exception(e)
logging.warning("Continuing with tests, gurobi not installed")
return
from remat.core.solvers.strategy_optimal_ilp import solve_ilp_gurobi
model = get_keras_model("test")
g = dfgraph_from_keras(mod=model)
assert g.size_fwd == 6
budget = sum(g.cost_ram.values()) + g.cost_ram_parameters
scheduler_result = solve_ilp_gurobi(g, budget)
assert scheduler_result.feasible
assert scheduler_result.schedule_aux_data.cpu <= sum(g.cost_cpu.values())
assert scheduler_result.schedule_aux_data.activation_ram <= sum(
g.cost_cpu.values())
assert scheduler_result.schedule_aux_data.peak_ram <= budget
from remat.core.solvers.strategy_checkpoint_all import solve_checkpoint_all, solve_checkpoint_all_ap
from remat.core.solvers.strategy_checkpoint_last import solve_checkpoint_last_node
from remat.core.solvers.strategy_chen import solve_chen_greedy, solve_chen_sqrtn
from remat.core.solvers.strategy_griewank import solve_griewank
from experiments.common.load_keras_model import get_keras_model
from remat.core.solvers.strategy_checkpoint_all import solve_checkpoint_all
from remat.tensorflow2.extraction import dfgraph_from_keras
if __name__ == "__main__":
model = get_keras_model("test")
g = dfgraph_from_keras(mod=model)
scheduler_result = solve_checkpoint_all(g)
print(scheduler_result.schedule)
cost_model = None
else:
cost_model = CostModel(model_name, args.platform, log_base, quantization=5)
cost_model.fit()
# gen redis key
if cost_model is None:
key_list = ["flops", args.batch_size]
else:
key_list = [cost_model.platform, cost_model.quantization, args.batch_size]
redis_cost_key = "_".join(map(str, key_list))
# load model from Keras
logger.info(f"Loading model {model_name}")
model = get_keras_model(model_name, input_shape=args.input_shape)
g = dfgraph_from_keras(model, batch_size=args.batch_size, cost_model=cost_model,
loss_cpu_cost=0, loss_ram_cost=(4 * args.batch_size))
result_dict = pickle.load((log_base / 'result_dict.pickle').open('rb'))
simrd_eval_points = pickle.load((log_base / 'simrd_eval_points.pickle').open('rb'))
simrd_results = []
for heuristic in SIMRD_HEURISTICS:
simrd_results.append(run_simrd(g, heuristic, simrd_eval_points, SIMRD_LIVENESS))
# save simrd results and heuristics used
pickle.dump(simrd_results, (log_base / 'simrd_results.pickle').open('wb'), \
protocol=pickle.HIGHEST_PROTOCOL)
pickle.dump(SIMRD_HEURISTICS, (log_base / 'simrd_heuristics.pickle').open('wb'), \
protocol=pickle.HIGHEST_PROTOCOL)
bs_fwd2xcost: Dict[int, int] = {}
rg = list(
range(args.batch_size_min, args.batch_size_max,
args.batch_size_increment))
for bs in tqdm(rg, desc="Event dispatch"):
while not ray.is_initialized():
ray.init(temp_dir="/tmp/ray_checkpoint_" +
str(str(uuid.uuid4())[:8]),
redis_password=str(uuid.uuid1()),
num_cpus=os.cpu_count() - 2)
futures = []
# load model at batch size
g = dfgraph_from_keras(model,
batch_size=bs,
cost_model=cost_model,
loss_cpu_cost=0,
loss_ram_cost=(4 * bs))
bs_fwd2xcost[bs] = sum(g.cost_cpu_fwd.values()) + sum(
g.cost_cpu.values())
bs_param_ram_cost[bs] = g.cost_ram_fixed
render_dfgraph(g, log_base, name=model_name)
# run constant baselines
result_dict[bs][SolveStrategy.CHEN_SQRTN_NOAP] = [
solve_chen_sqrtn(g, False)
]
futures.extend([
ray.remote(num_cpus=1)(solve_checkpoint_all).remote(g),
ray.remote(num_cpus=1)(solve_checkpoint_all_ap).remote(g),
ray.remote(num_cpus=1)(solve_checkpoint_last_node).remote(g),