philipperemy/keras-tcn

Question: TCN Dilated Convolutions?

asapsmc opened this issue · 3 comments

First of all, sorry for filling an issue, but I had no option to just make a question.
My question is simple, how did you draw the dilated convolutions? Was it with Tikz? (I have to do some diagrams including TCN's but I'm unable to achieve so perfect renditions of the dilated convolutions in tikz.
Thanks in advance.

krzim commented

It's almost definitely networkx. I recreated some of them using this code but I can't make claims for it's generality or accuracy since I just ripped it out of an old notebook and pasted it here. Hopefully it helps you out!

import networkx as nx
import numpy as np
from matplotlib import pyplot as plt

block_factor = 1  # assuming only one convolution in a block for simplicity
dilation_factor = 4
num_blocks = 3
num_stacks = 1
kernel_size = 4
sum_dilations = sum(dilation_factor**i for i in range(num_blocks))
receptive_field = 1 + block_factor * (kernel_size - 1) * num_stacks * sum_dilations 

fill_inactive = False
if fill_inactive:
    graph_width = 3 * receptive_field
else:
    graph_width = receptive_field

# intialize the graphs with white nodes
active_graph = nx.create_empty_copy(nx.grid_2d_graph(num_blocks + 1, graph_width))
inactive_graph = nx.create_empty_copy(active_graph)
nx.set_node_attributes(active_graph, values="white", name="color")
nx.set_node_attributes(inactive_graph, values="white", name="color")


# color input nodes
input_nodes = [node for node in active_graph.nodes if node[0] == 0]
nx.set_node_attributes(active_graph, {node: "gold" for node in input_nodes}, name="color")

# color output nodes
output_nodes = [node for node in active_graph.nodes if node[0] == num_blocks]
nx.set_node_attributes(active_graph, {node: "red" for node in output_nodes}, name="color")

# highlight active nodes
for block_num in range(num_blocks, 0, -1):
    dilation = dilation_factor ** block_num
    next_dilation = dilation_factor ** (block_num - 1)
    for i in range(graph_width - 1, (next_dilation * (kernel_size - 1)) - 1, -dilation):
        nx.set_node_attributes(active_graph, {(block_num, i): {"color": "deepskyblue"}})
        for j in range(kernel_size):
            active_graph.add_edge((block_num - 1, i - (j * next_dilation)), (block_num, i), style="solid")

# fill in inactive nodes
if fill_inactive:
    for block_num in range(num_blocks):
        dilation = dilation_factor ** block_num
        for i in range(graph_width):
            inactive_graph.add_edge((block_num, i), (block_num + 1, i), style="dashed")
            if i > dilation:
                inactive_graph.add_edge((block_num, i-dilation), (block_num + 1, i), style="dashed")

plt.figure(figsize=(receptive_field, 0.3 * receptive_field))
pos = {(x,y):(y,x) for x,y in active_graph.nodes()}
nx.draw(
    active_graph,
    pos=pos,
    arrows=True,
    arrowstyle="-|>",
    arrowsize=10,
    width=1.5,
    style="solid",
    node_color=nx.get_node_attributes(active_graph, "color").values(),
    edgecolors="black",
)
nx.draw_networkx_edges(
    inactive_graph,
    pos=pos,
    arrows=True,
    arrowstyle="-|>",
    arrowsize=8,
    width=0.8,
    style="dashed",
)
if fill_inactive:
    plt.xlim(left=receptive_field - 0.5, right=2 * receptive_field - 0.5)
plt.show()

Big +1 to @krzim for answering on that one!

Indeed, my fault for lack of aknwoledgement! Thanks @krzim