Question: TCN Dilated Convolutions? about keras-tcn HOT 3 CLOSED

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()