def connect(e_from, e_to): cs_from = stz.coords_from_bbox_with_fn(e_from, stz.top_center_coords) cs_to = stz.coords_from_bbox_with_fn(e_to, stz.bottom_center_coords) s_fmt = fmt.combine_tikz_strs( [fmt.line_width(line_width), fmt.arrow_heads("end")]) return stz.line_segment(cs_from, cs_to, s_fmt)
def connect_residual(e_with_addnorm, bottom_spacing, side_spacing): from_cs = stz.coords_from_bbox_with_fn(e_with_addnorm, stz.bottom_center_coords) from_cs = stz.translate_coords_vertically(from_cs, -bottom_spacing) if side_spacing < 0: to_cs = stz.coords_from_bbox_with_fn(e_with_addnorm[0], stz.left_center_coords) else: to_cs = stz.coords_from_bbox_with_fn(e_with_addnorm[0], stz.right_center_coords) cs = stz.translate_coords_horizontally(from_cs, side_spacing) return stz.open_path([from_cs, cs, [cs[0], to_cs[1]], to_cs], s_con)
def independent_hyperparameter(name, values_expr, value=None): e = stz.ellipse([0, 0], h_width / 2.0, h_height / 2.0, hyperp_s_fmt) l = stz.latex(stz.center_coords(e), "\\textbf{%s}" % name) fn_cs = stz.coords_from_bbox_with_fn(e, stz.right_center_coords) if value is None: l_vs = stz.latex(fn_cs, "\\textbf{[%s]}" % (values_expr, ), unassigned_h_s_fmt) return [e, l, l_vs] else: v_cs = stz.coords_from_bbox_with_fn(e, stz.right_center_coords) l_v = stz.latex(v_cs, "\\textbf{%s}" % value, assigned_h_s_fmt) return [e, l, l_v]
def dependent_hyperparameter(name, hyperp_names, fn_expr, value=None): e = stz.ellipse([0, 0], h_width / 2.0, h_height / 2.0, hyperp_s_fmt) if value is None: e["horizontal_radius"] *= 2.1 * e["horizontal_radius"] l_cs = stz.center_coords(e) if value is None: l_cs = stz.translate_coords_horizontally(l_cs, 0.1) l = stz.latex(l_cs, "\\textbf{%s}" % name) if value is None: fn_cs = stz.coords_from_bbox_with_fn(e, stz.right_center_coords) l_fn = stz.latex(fn_cs, "\\textbf{fn: %s}" % (fn_expr, ), unassigned_h_s_fmt) p = property("x", 0.25, 0.7) p_cs = stz.translate_coords_horizontally( stz.coords_from_bbox_with_fn(e, stz.left_center_coords), 0.1) stz.translate_bbox_left_center_to_coords(p, p_cs) return [e, l, l_fn, p] else: v_cs = stz.coords_from_bbox_with_fn(e, stz.right_center_coords) l_v = stz.latex(v_cs, "\\textbf{%s}" % value, assigned_h_s_fmt) return [e, l, l_v]
def bar_plot(data, label_strs): cs = [0, 0] rs = [ stz.rectangle_from_width_and_height(cs, rectangle_height, x, r_fmt) for x in data ] stz.distribute_vertically_with_spacing(rs, rectangle_spacing) frame_height = (len(data) * rectangle_height + (len(data) - 1) * rectangle_spacing + 2.0 * margin) frame = stz.rectangle_from_width_and_height(cs, frame_height, frame_width) stz.align_centers_vertically([frame, rs], 0.0) labels = [] for i, r in enumerate(rs): l_cs = stz.coords_from_bbox_with_fn(r, stz.left_center_coords) l_cs = stz.translate_coords_horizontally(l_cs, -label_spacing) l_s = label_strs[i] lab = stz.latex(l_cs, l_s, fmt.anchor('right_center')) labels.append(lab) ticks = get_ticks(frame) return [rs, labels, frame, ticks]
def connect_horizontally(e_from, e_to): s_fmt = fmt.arrow_heads("end") cs_from = stz.coords_from_bbox_with_fn(e_from, stz.right_center_coords) cs_to = stz.coords_from_bbox_with_fn(e_to, stz.left_center_coords) return stz.line_segment(cs_from, cs_to, s_fmt)
def arrow_out(e): from_cs = stz.coords_from_bbox_with_fn(e, stz.top_center_coords) to_cs = stz.translate_coords_vertically(from_cs, arrow_length) return stz.line_segment(from_cs, to_cs, s_arr)
def arrow_in(e): to_cs = stz.coords_from_bbox_with_fn(e, stz.bottom_center_coords) from_cs = stz.translate_coords_vertically(to_cs, -arrow_length) return stz.line_segment(from_cs, to_cs, s_arr)
def connect_straight_with_arrow(e_from, e_to): from_cs = stz.coords_from_bbox_with_fn(e_from, stz.top_center_coords) to_cs = stz.coords_from_bbox_with_fn(e_to, stz.bottom_center_coords) return stz.line_segment(from_cs, to_cs, s_arr)
def connect_straight_horizontal(e_from, e_to): from_cs = stz.coords_from_bbox_with_fn(e_from, stz.right_center_coords) to_cs = stz.coords_from_bbox_with_fn(e_to, stz.left_center_coords) return stz.line_segment(from_cs, to_cs, s_lw)
def frame(frame_idx): assert frame_idx >= 0 and frame_idx <= 3 c1 = conv2d(1) o = optional(1) r1 = repeat(1) r2 = repeat(2) cc = concat(1) c2 = conv2d(2) c3 = conv2d(3) c4 = conv2d(4) d = dropout(1) stz.distribute_horizontally_with_spacing([r1, r2], horizontal_module_spacing) stz.distribute_horizontally_with_spacing([c2, [c3, c4]], horizontal_module_spacing) modules = [] if frame_idx == 0: stz.distribute_vertically_with_spacing([cc, [r1, r2], o, c1], vertical_module_spacing) stz.align_centers_horizontally([cc, [r1, r2], o, c1], 0) modules.extend([c1, o, r1, r2, cc]) else: stz.distribute_vertically_with_spacing([c4, c3], vertical_module_spacing) stz.distribute_horizontally_with_spacing([c2, [c3, c4]], horizontal_module_spacing) stz.align_centers_vertically([[c3, c4], c2], 0) if frame_idx == 1: stz.distribute_vertically_with_spacing([cc, [c2, c3, c4], o, c1], vertical_module_spacing) stz.align_centers_horizontally([cc, [c2, c3, c4], o, c1], 0) modules.extend([c1, o, c2, c3, c4, cc]) else: stz.distribute_vertically_with_spacing([cc, [c2, c3, c4], d, c1], vertical_module_spacing) stz.align_centers_horizontally([cc, [c2, c3, c4], d, c1], 0) modules.extend([c1, d, c2, c3, c4, cc]) module_connections = [] if frame_idx == 0: module_connections.extend([ connect_modules(c1, o, 0, 0), connect_modules(o, r1, 0, 0), connect_modules(o, r2, 0, 0), connect_modules(r1, cc, 0, 0), connect_modules(r2, cc, 0, 1), ]) else: if frame_idx == 1: module_connections.extend([ connect_modules(c1, o, 0, 0), connect_modules(o, c2, 0, 0), connect_modules(o, c3, 0, 0), ]) else: module_connections.extend([ connect_modules(c1, d, 0, 0), connect_modules(d, c2, 0, 0), connect_modules(d, c3, 0, 0), ]) module_connections.extend([ connect_modules(c3, c4, 0, 0), connect_modules(c2, cc, 0, 0), connect_modules(c4, cc, 0, 1), ]) # # hyperparameters if frame_idx <= 1: h_o = independent_hyperparameter("IH-2", "0, 1") else: h_o = independent_hyperparameter("IH-2", "0, 1", "1") if frame_idx <= 0: h_r1 = dependent_hyperparameter("DH-1", ["x"], "2*x") h_r2 = independent_hyperparameter("IH-3", "1, 2, 4") else: h_r1 = dependent_hyperparameter("DH-1", ["x"], "2*x", "2") h_r2 = independent_hyperparameter("IH-3", "1, 2, 4", "1") if frame_idx <= 2: h_c1 = independent_hyperparameter("IH-1", "64, 128") h_c2 = independent_hyperparameter("IH-4", "64, 128") h_c3 = independent_hyperparameter("IH-5", "64, 128") h_c4 = independent_hyperparameter("IH-6", "64, 128") h_d = independent_hyperparameter("IH-7", "0.25, 0.5") else: h_c1 = independent_hyperparameter("IH-1", "64, 128", "64") h_c2 = independent_hyperparameter("IH-4", "64, 128", "128") h_c3 = independent_hyperparameter("IH-5", "64, 128", "128") h_c4 = independent_hyperparameter("IH-6", "64, 128", "64") h_d = independent_hyperparameter("IH-7", "0.25, 0.5", "0.5") def place_hyperp_right_of(h, m): y_p = stz.center_coords(m[3])[1] stz.align_centers_vertically([h], y_p) stz.place_to_the_right(h, m, spacing_between_module_and_hyperp) hyperparameters = [] place_hyperp_right_of(h_c1, c1) if frame_idx in [0, 1]: place_hyperp_right_of(h_o, o) hyperparameters.append(h_o) if frame_idx == 0: place_hyperp_right_of(h_r1, r2) stz.place_above_and_align_to_the_right(h_r2, h_r1, 0.8) hyperparameters.extend([h_r1, h_r2, h_c1]) else: place_hyperp_right_of(h_c1, c1) place_hyperp_right_of(h_c3, c3) place_hyperp_right_of(h_c4, c4) stz.place_below(h_c2, h_c1, 3.0) hyperparameters.extend([h_c1, h_c2, h_c3, h_c4]) if frame_idx in [2, 3]: place_hyperp_right_of(h_d, d) hyperparameters.extend([h_d]) unreachable_hyperps = [] if frame_idx == 1: stz.distribute_vertically_with_spacing([h_r1, h_r2], 0.2) unreachable_hyperps.extend([h_r1, h_r2]) if frame_idx >= 2: stz.distribute_vertically_with_spacing([h_o, h_r1, h_r2], 0.2) unreachable_hyperps.extend([h_r1, h_r2, h_o]) hyperparameters.extend(unreachable_hyperps) cs_fn = lambda e: stz.coords_from_bbox_with_fn(e, stz.left_center_coords) if frame_idx == 0: stz.translate_bbox_left_center_to_coords(h_r2, cs_fn([h_o, h_r1])) elif frame_idx == 1: stz.translate_bbox_left_center_to_coords(h_c2, cs_fn([h_o, h_c3])) else: stz.translate_bbox_left_center_to_coords(h_c2, cs_fn([h_d, h_c3])) hyperp_connections = [ connect_hyperp_to_module(h_c1, c1, 0), ] if frame_idx in [0, 1]: hyperp_connections.extend([connect_hyperp_to_module(h_o, o, 0)]) if frame_idx == 0: hyperp_connections.extend([ connect_hyperp_to_module(h_r1, r2, 0), connect_hyperp_to_module(h_r2, r1, 0), connect_hyperp_to_hyperp(h_r2, h_r1) ]) else: hyperp_connections.extend([ connect_hyperp_to_module(h_c2, c2, 0), connect_hyperp_to_module(h_c3, c3, 0), connect_hyperp_to_module(h_c4, c4, 0), ]) if frame_idx in [2, 3]: hyperp_connections.append(connect_hyperp_to_module(h_d, d, 0)) f = stz.rectangle_from_width_and_height([0, 0], frame_height, frame_width, frame_s_fmt) e = [modules, module_connections, hyperparameters, hyperp_connections] stz.translate_bbox_center_to_coords( f, stz.translate_coords_horizontally(stz.center_coords(e), 0.8)) if len(unreachable_hyperps) > 0: stz.translate_bbox_bottom_right_to_coords(unreachable_hyperps, stz.bbox(e)[1]) # frame id s = ["a", "b", "c", "d"][frame_idx] label = [stz.latex([0, 0], "\\Huge \\textbf %s" % s)] stz.translate_bbox_top_left_to_coords( label, stz.translate_coords_antidiagonally( stz.coords_from_bbox_with_fn(f, stz.top_left_coords), 0.6)) return e + [f, label]
def connect_hyperp_to_hyperp(h_from, h_to): return stz.line_segment( stz.coords_from_bbox_with_fn(h_from[:2], stz.right_center_coords), stz.coords_from_bbox_with_fn(h_to[3], stz.top_center_coords), connect_s_fmt)
def connect_hyperp_to_module(h, m, property_idx): return stz.line_segment( stz.coords_from_bbox_with_fn(h[:2], stz.left_center_coords), stz.coords_from_bbox_with_fn(m[3][property_idx], stz.right_center_coords), connect_s_fmt)
def connect_modules(m_from, m_to, output_idx, input_idx): return stz.line_segment( stz.coords_from_bbox_with_fn(m_from[2][output_idx], stz.bottom_center_coords), stz.coords_from_bbox_with_fn(m_to[1][input_idx], stz.top_center_coords), connect_s_fmt)
def label_left(e, expr): cs = stz.coords_from_bbox_with_fn(e, stz.left_center_coords) cs = stz.translate_coords(cs, -label_spacing, 0.1) return stz.latex(cs, "\\scriptsize{%s}" % expr)
def label_below(e, expr): cs_ref = stz.coords_from_bbox_with_fn(e, stz.bottom_center_coords) cs = stz.translate_coords_vertically(cs_ref, -label_spacing) return stz.latex(cs, expr)
def label_right(e, expr): cs_ref = stz.coords_from_bbox_with_fn(e, stz.right_center_coords) cs = stz.translate_coords_horizontally(cs_ref, label_spacing) return stz.latex(cs, expr)
connections.extend([e1, e2]) tri_cs, translated_tri_cs = trident_coords(mmha_an) e1 = stz.open_path([ translated_tri_cs[1], stz.bottom_left_coords(translated_tri_cs[1], tri_cs[0]), tri_cs[0] ], s_con) e2 = stz.open_path([ translated_tri_cs[1], stz.bottom_right_coords(translated_tri_cs[1], tri_cs[2]), tri_cs[2] ], s_con) connections.extend([e1, e2]) tri_cs, translated_tri_cs = trident_coords(mha2_an) # the leftmost two from_cs = stz.coords_from_bbox_with_fn(ff1_an, stz.top_center_coords) mid_cs = stz.midway_coords(stz.bbox(bb1)[1], stz.bbox(bb2)[0]) to_cs = translated_tri_cs[0] cs = stz.translate_coords_vertically(from_cs, 1.0) e1 = stz.open_path([ from_cs, cs, [mid_cs[0], cs[1]], [mid_cs[0], to_cs[1]], translated_tri_cs[0], tri_cs[0] ], s_con) e2 = stz.open_path([ from_cs, cs, [mid_cs[0], cs[1]], [mid_cs[0], to_cs[1]], translated_tri_cs[1], tri_cs[1] ], s_con) # the rightmost one from_cs = stz.coords_from_bbox_with_fn(mmha_an, stz.top_center_coords)