How to split relay expr in TVM

#1

Hi everyone: I have a question. Suppose I want to translate the following torch codes into relay front-end codes. I am confused about how to split relay expr into two parts. For example, here is my torch code:

            offset_mask = nn.Conv2d(
               in_channels, out_channels,
                kernel_size=3,
                padding=1)(data)
            offset = offset_mask[:, :18, :, :]
            mask = offset_mask[:, -9:, :, :].sigmoid()

I am trying to translate above codes into relay, but I am stuck how to translate the last two codes which require me to split relay.expr into two parts.

   offset_mask = relay.nn.conv2d(data, weight, strides=(1,1) padding = (1,1), kernel_size = (3,3))

After that, how am I gonna to do?

#2

Take a look at strided_slice (tvm.relay — tvm 0.8.dev0 documentation). I think in this case:

offset = relay.strided_slice(offset_mask, begin=(0, 0, 0, 0), end=(-1, 18, -1, -1), slice_mode="size")
# assuming NCHW is preserved, offset_mask shape is [N, out_channels, H, W]
mask = relay.strided_slice(offset_mask, begin=(0, out_channels - 9, 0, 0), end=(-1, -1, -1, -1), slice_mode="size")
mask = relay.sigmoid(mask)

Let me know if this helps

#3

I used relay.split function (tvm.relay — tvm 0.8.dev0 documentation), it works!