|
import json |
|
|
|
import numpy as np |
|
import drawsvg as draw |
|
import colorsys |
|
import tempfile, os |
|
|
|
def filter_time(e, start, end): |
|
|
|
|
|
return e.completion_time >= start and (end is None or e.completion_time <= end) |
|
|
|
|
|
def load_json_data(filename, start=0, end=None, time_scale=1): |
|
with open(filename) as f: |
|
data = json.loads(f.read()) |
|
fbw_types = {"F", "B", "W", "Optimizer"} |
|
return [[{ |
|
"type": e["type"], |
|
"start_time": int(max(e["start_time"] - start, 0)) * time_scale, |
|
"completion_time": int(e["completion_time"] - start) * time_scale, |
|
"minibatch": e.get("minibatch", None), |
|
} for e in dev_evs |
|
if e["type"] in fbw_types and filter_time(e, start, end) |
|
] for dev_evs in data] |
|
|
|
|
|
ENABLE_BORDER = True |
|
ENABLE_BATCH_ID = True |
|
ENABLE_EDGE_BLUR = False |
|
SCALE_FACTOR = 2 |
|
S = SCALE_FACTOR |
|
|
|
|
|
TIME_PER_UNIT = 4000 // SCALE_FACTOR |
|
|
|
|
|
def to_color_fmt(c): |
|
|
|
return f"#{hex(c[0])[2:]}{hex(c[1])[2:]}{hex(c[2])[2:]}" |
|
|
|
|
|
GREYSCALE_WEIGHTS = np.array([0.299, 0.587, 0.114]) |
|
|
|
|
|
def to_greyscale(color): |
|
c = np.dot(GREYSCALE_WEIGHTS, color[:3].astype(float)).astype(int) |
|
return np.array([c, c, c, 255]) |
|
|
|
|
|
COLOR_VALUE_MAP = { |
|
"F": np.array([57, 122, 242]), |
|
"B": np.array([62, 181, 191]), |
|
|
|
|
|
"W": np.array([41, 137, 64]), |
|
|
|
|
|
"Optimizer": np.array([255, 217, 102]), |
|
} |
|
|
|
|
|
COLOR_MAP = {k: to_color_fmt(v) for k, v in COLOR_VALUE_MAP.items()} |
|
|
|
|
|
|
|
BORDER_SIZE = 1 |
|
SPAN_HEIGHT = SCALE_FACTOR * 10 |
|
FONT_SIZE = SCALE_FACTOR * 10 |
|
TITLE_WIDTH = SCALE_FACTOR * 60 |
|
CENTER_TITLE_HEIGHT = SPAN_HEIGHT * 6 |
|
|
|
WHITE = to_color_fmt(np.array([255, 255, 255, 255])) |
|
BLACK = to_color_fmt(np.array([0, 0, 0, 255])) |
|
|
|
|
|
class DrawCtx: |
|
def __init__(self, d, oy, ox): |
|
assert not isinstance(d, DrawCtx) |
|
self.d = d |
|
self.oy = oy |
|
self.ox = ox |
|
|
|
@classmethod |
|
def from_base_ctx(cls, base_ctx, oy, ox): |
|
assert isinstance(base_ctx, DrawCtx) |
|
return cls(base_ctx.d, base_ctx.oy + oy, base_ctx.ox + ox) |
|
|
|
def width(self): |
|
return self.d.width |
|
|
|
def height(self): |
|
return self.d.height |
|
|
|
def line(self, sy, sx, ey, ex, width=None): |
|
self.d.append(draw.Line( |
|
self.ox + sx, |
|
self.oy + sy, |
|
self.ox + ex, |
|
self.oy + ey, |
|
stroke='black', |
|
stroke_width=width or BORDER_SIZE, |
|
)) |
|
|
|
def rect(self, sy, sx, h, w, color): |
|
self.d.append(draw.Rectangle( |
|
self.ox + sx, |
|
self.oy + sy, |
|
w, h, |
|
fill=color, |
|
shape_rendering="geometricPrecision", |
|
)) |
|
|
|
def rect_frame(self, sy, sx, h, w): |
|
self.d.append(draw.Rectangle( |
|
self.ox + sx, |
|
self.oy + sy, |
|
w, h, |
|
fill="none", |
|
stroke=BLACK, |
|
stroke_width=BORDER_SIZE, |
|
)) |
|
|
|
def text(self, y, x, text, anchor="middle", font_scale=1, fill='black'): |
|
font_size = FONT_SIZE * font_scale |
|
tl = len(text) * font_size // 2 |
|
self.d.append(draw.Text( |
|
text, font_size, |
|
self.ox + x, |
|
|
|
self.oy + y + font_size - 3, |
|
textLength=tl, lengthAdjust='spacing', |
|
text_anchor=anchor, |
|
font_family="Times New Roman", |
|
fill=fill, |
|
|
|
|
|
)) |
|
|
|
|
|
def change_color_sat(c, percentage): |
|
c = c.astype(float) / 255.0 |
|
(h, s, v) = colorsys.rgb_to_hsv(c[0], c[1], c[2]) |
|
s *= percentage |
|
r, g, b = colorsys.hsv_to_rgb(h, s, v) |
|
c = np.array([r, g, b]) * 255 |
|
return c.astype(int) |
|
|
|
|
|
def draw_experiment_and_schedule(exp_events, sched_events, output_filename, tail=10): |
|
exp_canvas_info = CanvasInfo(exp_events, tail, 0) |
|
sched_canvas_info = CanvasInfo(sched_events, tail, 0, False) |
|
width = max(exp_canvas_info.get_canvas_size()[1], sched_canvas_info.get_canvas_size()[1]) |
|
height = exp_canvas_info.get_canvas_size()[0] + sched_canvas_info.get_canvas_size()[0] |
|
|
|
include_w = True |
|
|
|
|
|
d = draw.Drawing(width, height, origin="top-left") |
|
ctx = DrawCtx(d, 0, 0) |
|
plot_events(ctx, sched_events, "", sched_canvas_info, include_w, include_o=False, include_info=False) |
|
|
|
|
|
|
|
|
|
|
|
exp_ctx = DrawCtx.from_base_ctx(ctx, sched_canvas_info.get_canvas_size()[0], 0) |
|
plot_events(exp_ctx, exp_events, "", exp_canvas_info, include_w, include_o=True) |
|
|
|
d.save_svg(output_filename) |
|
|
|
|
|
def draw_events(events, output_filename, include_w=True, include_o=True, tail=50, longest_time=None): |
|
canvas_info = CanvasInfo(events, tail, center_title_height=0, enable_info=True, longest_time=longest_time) |
|
max_len = canvas_info.max_len |
|
|
|
|
|
height, width = canvas_info.get_canvas_size() |
|
|
|
d = draw.Drawing(width, height, origin="top-left") |
|
ctx = DrawCtx(d, 0, 0) |
|
|
|
plot_events(ctx, events, "", canvas_info, include_w, include_o) |
|
d.save_svg(output_filename) |
|
|
|
|
|
class CanvasInfo: |
|
def __init__(self, events, tail, center_title_height=CENTER_TITLE_HEIGHT, enable_info=True, longest_time=None): |
|
|
|
last_time = max(max([e["completion_time"] for e in dev_evs]) for dev_evs in events) if longest_time is None else longest_time |
|
self.max_len = (last_time + TIME_PER_UNIT - 1) // TIME_PER_UNIT + tail |
|
|
|
self.height = SPAN_HEIGHT * len(events) + BORDER_SIZE * (len(events) + 1) |
|
color_text_row_height = int(SPAN_HEIGHT * 1.6) |
|
self.color_text_height = color_text_row_height + BORDER_SIZE |
|
self.info_height = SPAN_HEIGHT + color_text_row_height + 3 * BORDER_SIZE |
|
if not enable_info: |
|
self.info_height /= 2 |
|
self.center_title_height = center_title_height |
|
|
|
|
|
def get_canvas_size(self): |
|
|
|
return self.height + self.info_height + self.center_title_height, self.max_len + TITLE_WIDTH |
|
|
|
|
|
def plot_events(ctx, events, title_text: str, canvas_info: CanvasInfo, include_w=True, include_o=True, include_info=True): |
|
max_len = canvas_info.max_len |
|
height = canvas_info.height |
|
color_text_height = canvas_info.color_text_height |
|
info_height = canvas_info.info_height |
|
|
|
data_ctx = DrawCtx.from_base_ctx(ctx, 0, TITLE_WIDTH) |
|
|
|
for i, evs in enumerate(events): |
|
h = i * SPAN_HEIGHT + (i + 1) * BORDER_SIZE |
|
for e in evs: |
|
start = BORDER_SIZE + e["start_time"] // TIME_PER_UNIT |
|
end = BORDER_SIZE + e["completion_time"] // TIME_PER_UNIT |
|
if start == end or not ENABLE_EDGE_BLUR: |
|
plot_span(data_ctx, start, end, h, COLOR_MAP[e["type"]]) |
|
else: |
|
plot_span(data_ctx, start + 1, end - 1, h, COLOR_MAP[e["type"]]) |
|
|
|
c = change_color_sat( |
|
COLOR_VALUE_MAP[e["type"]], |
|
(e["start_time"] / TIME_PER_UNIT) % 1.0) |
|
plot_span(data_ctx, start, start + 1, h, to_color_fmt(c)) |
|
c = change_color_sat( |
|
COLOR_VALUE_MAP[e["type"]], |
|
(e["completion_time"] / TIME_PER_UNIT) % 1.0) |
|
plot_span(data_ctx, end - 1, end, h, to_color_fmt(c)) |
|
|
|
if ENABLE_BATCH_ID: |
|
minibatch = str(e["minibatch"]) |
|
center = (start + end) // 2 |
|
data_ctx.text(h, center, minibatch, font_scale=0.6, fill='black' if e["chunk"] == 0 else 'white') |
|
if ENABLE_BORDER: |
|
data_ctx.line(h+SPAN_HEIGHT, 0, h+SPAN_HEIGHT+BORDER_SIZE, max_len - 1) |
|
|
|
if ENABLE_BORDER: |
|
data_ctx.line(0, 0, 0, max_len - 1) |
|
data_ctx.line(0, 0, height, 0) |
|
data_ctx.line(0, max_len - 1, height, max_len - 1) |
|
|
|
dev_title_ctx = DrawCtx.from_base_ctx(ctx, 0, 0) |
|
ndev = len(events) |
|
add_devices(dev_title_ctx, ndev) |
|
|
|
if not include_info: |
|
return |
|
|
|
info_height = ndev * SPAN_HEIGHT + (ndev + 1) * BORDER_SIZE |
|
info_ctx = DrawCtx.from_base_ctx(ctx, info_height, 0) |
|
add_info(info_ctx, color_text_height, include_w, include_o) |
|
|
|
if title_text: |
|
center_title_ctx = DrawCtx.from_base_ctx(info_ctx, canvas_info.info_height, 0) |
|
add_center_title(center_title_ctx, title_text) |
|
|
|
|
|
def plot_span(ctx, start, end, h, color, ): |
|
ctx.rect(h, start, SPAN_HEIGHT, end - start, color) |
|
if ENABLE_BORDER: |
|
ctx.rect_frame(h-BORDER_SIZE, start, SPAN_HEIGHT + BORDER_SIZE, end - start) |
|
|
|
|
|
def add_devices(ctx, devs): |
|
for i in range(devs): |
|
h = i * SPAN_HEIGHT + (i + 1) * BORDER_SIZE |
|
ctx.text(h, 6 * SCALE_FACTOR, "Device {}".format(i), "left") |
|
|
|
|
|
def add_info(ctx, color_text_height, include_w=True, include_o=True): |
|
div = 4 + int(include_w) + int(include_o) |
|
f_start = ctx.width() // div |
|
b_start = ctx.width() // div * 2 |
|
w_start = ctx.width() // div * 3 |
|
o_start = ctx.width() // div * 4 |
|
|
|
block_w = 25 * SCALE_FACTOR |
|
plot_span(ctx, f_start, f_start+block_w, color_text_height + BORDER_SIZE, COLOR_MAP["F"]) |
|
plot_span(ctx, b_start, b_start+block_w, color_text_height + BORDER_SIZE, COLOR_MAP["B"]) |
|
if include_w: |
|
plot_span(ctx, w_start, w_start+block_w, color_text_height + BORDER_SIZE, COLOR_MAP["W"]) |
|
if include_o: |
|
plot_span(ctx, o_start, o_start+block_w, color_text_height + BORDER_SIZE, COLOR_MAP["Optimizer"]) |
|
|
|
ctx.text(0, 6 * SCALE_FACTOR, "Time", "left") |
|
draw_arrow(ctx, SPAN_HEIGHT // 2 + BORDER_SIZE + 1, 65 * SCALE_FACTOR, 50 * SCALE_FACTOR) |
|
|
|
block_w = 30 * SCALE_FACTOR |
|
ctx.text(color_text_height, f_start + block_w, "F", "left") |
|
ctx.text(color_text_height, b_start + block_w, |
|
"B", "left") |
|
if include_w: |
|
ctx.text(color_text_height, w_start + block_w, "W", "left") |
|
if include_o: |
|
ctx.text(color_text_height, o_start + block_w, "Optimizer Step", "left") |
|
|
|
|
|
def add_center_title(ctx: DrawCtx, text): |
|
ctx.text(CENTER_TITLE_HEIGHT / 4, ctx.width() / 2, |
|
text, "middle", 2) |
|
|
|
|
|
def draw_arrow(ctx: DrawCtx, start_y, start_x, width, thickness=2): |
|
b = thickness * (SCALE_FACTOR // 2) |
|
ctx.line(start_y, start_x, start_y, start_x + width, b) |
|
ctx.line(start_y, start_x + width, start_y - 3*b, start_x + width - 3*b) |
|
ctx.line(start_y, start_x + width, start_y + 3*b, start_x + width - 3*b) |
|
|
|
|
|
def render_manual_graph(data, longest_time, enable_batch_id = False): |
|
global ENABLE_BORDER |
|
global ENABLE_BATCH_ID |
|
ENABLE_BORDER = True |
|
ENABLE_BATCH_ID = enable_batch_id |
|
fbw_types = {"F", "B", "W", "Optimizer"} |
|
start = 0 |
|
end = None |
|
time_scale= 1024 / longest_time * TIME_PER_UNIT |
|
events = [[{ |
|
"type": e.type, |
|
"start_time": int(max(e.start_time - start, 0)) * time_scale, |
|
"completion_time": int(e.completion_time - start) * time_scale, |
|
"minibatch": e.minibatch, |
|
"chunk": e.chunk if hasattr(e, "chunk") else 0, |
|
} for e in dev_evs |
|
if e.type in fbw_types and filter_time(e, start, end) |
|
] for dev_evs in data] |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
path = os.path.join(tempfile.mkdtemp(), 'a.svg') |
|
draw_events(events, path, include_w=True, include_o=False, tail=50, longest_time=longest_time * time_scale) |
|
return path |
|
|
|
|
|
def render_experiment_graph(): |
|
global ENABLE_BORDER |
|
global ENABLE_BATCH_ID |
|
global TIME_PER_UNIT |
|
ENABLE_BORDER = False |
|
ENABLE_BATCH_ID = False |
|
TIME_PER_UNIT = 200 // SCALE_FACTOR |
|
TIME_PER_UNIT *= 12000 |
|
start_time = 1100000000 + 10000000 |
|
|
|
iter_time = 1290000000 |
|
end_time = start_time + iter_time |
|
exp_events = load_json_data("20-09-zero/zero-events.json", start_time, end_time) |
|
|
|
sched_events = load_json_data("schedule.json", time_scale=1000) |
|
draw_experiment_and_schedule(exp_events, sched_events, "pics/exp.svg") |
|
|
|
|
|
|
|
|
|
|
|
|
