Support for NHWC?

[nielsdekoeijer]

Hi, I stumbled across your repository recently as I also embarked on a similar journey in implementing tiny-cnn for audio ML problems. Your implementation mirrors mine so it's good to see someone else taking a similar approach! Great job.

One interesting result I came across was that using the NHWC layout (rather than the NCHW layout used in torch and onnxruntime) gave me a speedup. I got a benefit for both "direct" and "im2col" implementations.

Honestly, I'm just quite curious if this is something you've ever considered? (or better yet: if you've tried and have similar results?)

(p.s. I realize creating an issue may be a strange way to reach out, let me know if you prefer to delete this issue and communicate in another way!).

[Avafly]

Hi, it's totally fine to have the discussion in the issue~

I have not tried NHWC layout yet, but let me analyze the impact before implementation. I want to focus on the convolutional layer because it takes up the majority of the time. Since the direct way is too slow for my device, I just analyze the im2col approach.

Convolutional layer computation can be divided into two steps: im2col and GEMM. Because changing data layout shouldn't change the output shape of im2col, GEMM likely won't change much. Even so, GEMM still can reach peak CPU performance with different matrix shape (I've tried to optimize GEMM in another repo). So let's focus on im2col. In fact, in small models such as tiny-cnn, the elapsed time of im2col is often greater than the computation itself (i.e. GEMM). Thus I think im2col may be the key factor impacting speed.

The NHWC layout is more focused on local data (all pixel values of a region are kept in adjacent memory), which is aligns well with the convolution operation (it also computes region-by-region), whereas the NCHW layout has to jump across the channel space to fetch data from different channels in a region. So I believe NHWC layout would be more efficient.

However, I'm not certain how much this improvement would be in practice. Since the model is very tiny, data to be used may already be loaded in cache, so factors like layout might have less impact. For example, the maxpooling layer used in tiny-cnn is a naive implementation, but it's faster than the Caffe version. If the model is not tiny, inference frameworks like ncnn and OnnxRuntime are better to use since they use additional optimization methods like Winograd.

I'm very interested in trying a NHWC implementation, but I may not get to it right away because I have some other works to do. I'll let you know if there is any update.

[nielsdekoeijer]

If you'll entertain me, I can give you an example for now until I finalize my notes. My code seems to pass my unit tests but there are some warts that I'd like to get rid of first. I'm also not entirely sure if I'm implementing im2col or im2row to be honest. All this is a bit fuzzy in my mind still!

However, to give you the gist of how I've done it, without any guarantee that this is optimal, the loops I use are as follows:

Loops:

// input: (iN, iH, iW, iC)
// weight: (wM, wH, wW, wC)
// out: (oN, oH, oW, oC)
for (0..oH) |oh| { // for each output height
    for (0..oW) |ow| { // for each output width
        for (0..kH) |kh| { // for each kernel height
            if (0 <= ih and ih < iH) { // assert bounds height
                for (0..kW) |kw| { // for each kernel width
                    if (0 <= iw and iw < iW) { // assert bounds width
                        // memcpy scratch matrix [dst..dst + iC] from input [src .. src + iC]
                        // note: contiguous !! so that means good I guess
                    }
                }
            }
        }
    }
}

Let me know if it makes sense :-)

I'll publish the repo here once I've confirmed all is well.

[Avafly]

No matter im2col or im2row, they are both transforming the input tensor to a certain form for GEMM.

I think your loops make sense, and I've successfully made a simple Python code to implement a convolutional layer in this way. Adopting the NHWC layout in my existing code seems tricky though cause it isn't compatible in many parts (not just layout). I need to write a new cnn in NHWC layout from scratch.

I'll share updates here as I make progress.

[nielsdekoeijer]

Sounds great. I have now put my repository online here: https://github.com/nielsdekoeijer/research-efficient-conv2d

The code is not a work of art, but at least it passes my tests. Inside the repo, you can also see that I get a speedup with NHWC. This is because of the contiguous memory accesses, allowing for some SIMD stuff! Hopefully, this is due to merits of the layout and not just a mistake in my NCHW implementation. Regardless I'm curious what you come up with 👍.

[nielsdekoeijer]

Also, given your interest in implementing a nice gemm, perhaps you're aware but this MOOC is really nice:
https://www.cs.utexas.edu/~flame/laff/pfhp/week0-welcome.html

In addition, I agree with you that the im2col can take some time. Are you familiar with indirect im2col? A resource:
https://arxiv.org/pdf/1907.02129

I'm quite curious if indirect im2col solves some of the issues, but it essentially requires you write your own gemm (using the indirection buffers as described in the paper). That leaves the problem that writing a fast gemm is non-trivial (consider the first link as my first venture there).

[Avafly]

I implemented the conv2d layer under NCHW and NHWC layouts and compared them. Code is available at: https://github.com/Avafly/conv-layout-study.

In short, test results show NHWC was better in most cases for im2col and the conv2d layer.

The README mainly shows im2col's execution times for k=3/s=1/p=0. I actually tested other kernel sizes/strides/paddings, but didn't include them for brevity. Clone and run the repo if you're interested :D

Back to the question: Will tiny-cnn support NHWC? Considering that for 8x16x16 inputs, the speeds of two layouts are almost the same, and blobs in tiny-cnn are even smaller, I don't think NHWC layout would make much difference for tiny-cnn. Additionally, tiny-cnn is already very fast (about 0.009ms per image recognition). So my answer is I won't implement it for now.

Still thanks for making me think about how layouts affect model inference 👍. Please let me know if you have any idea about better im2col implementations or other things~