Relay Level Tiling of Conv2d or any operator

Questions
#1

Dear All,

I am wondering how can I write a Relay pass that tiles conv2d by the output channels (data partitioning) in Relay graph level?

For example, let us assume that I have some relay program like below, and I want to able to traverse the relay graph that contains this conv2d, and able to re-write it to one below with two conv2d?

kernel_size = 3
# N, C, H, W
input_shape = (1, 4, 10, 10)

# OIHW 
out_channels = 6
weight_shape = (**out_channels**, 4, kernel_size, kernel_size)

conv2d  = relay.nn.conv2d(
input,
weight,
channels = **out_channels **,
kernel_size = 3
)

Translate this into

conv2d_1  = relay.nn.conv2d(
input,
weight,
channels = **out_channels/2 **,
kernel_size = 3
)

conv2d_2  = relay.nn.conv2d(
input,
weight,
channels = **out_channels/2 **,
kernel_size = 3
)

conv2d_final = tvm.relay.concatenate (conv2d_1, conv2d_2 )
#2

You can write a small pass in Python to partition the program in the way that you would like. For example this below pass splits the inputs, and then adds, and concatenates them. You can employ a similar technique to do the partitioning that you want.

1 Like