tsinghua-ideal / Canvas

Canvas: End-to-End Kernel Architecture Search in Neural Networks
12 stars 4 forks source link

How do I condition CANVAS on runtime? #3

Open 5thGenDev opened 5 months ago

5thGenDev commented 5 months ago

I referred to this line here in CANVAS paper: "Then Canvas evaluates the candidate kernels using two classes of user-specified constraints and optimization goals, including those that can be analytically modeled and calculated (e.g., numbers of operations and parameters), and those that need experimental measurements (e.g., runtime and accuracy)"

LyricZhao commented 2 months ago

@5thGenDev Hi, very thankful to see you are trying this work and sorry for my late reply. For release consideration, I delete all my experiment code, but you can refer to this commit, which contains my experiment: https://github.com/tsinghua-ideal/Canvas/tree/30eb2d3bbe038c969bd8cff7580a7a672d705249. As shown in https://github.com/tsinghua-ideal/Canvas/blob/30eb2d3bbe038c969bd8cff7580a7a672d705249/experiments/search.py, after sampling a kernel by Canvas, we can count the FLOPs/params and filter the kernel. For runtime (model inference latency), the experiment in the paper uses a separate machine to TVMize kernels and evaluate the performance, this part is not fully automatic and not shown in the open-source code.

LyricZhao commented 2 months ago

Feel free to open/close this issue, if you need more code/help about the paper (e.g. TVM code generation and runtime evaluation), I am glad to provide :)

5thGenDev commented 2 months ago

Feel free to open/close this issue, if you need more code/help about the paper (e.g. TVM code generation and runtime evaluation), I am glad to provide :)

I have an idea to replace conv 3x3 layer on pretrained U-Net diffusion. Im not sure how I can somehow hook .sample() class from Canvas to replace conv 3x3 in different encoders/decoders block

LyricZhao commented 2 months ago

You could follow the example in README. Or for an existing network, you can iterate all modules, match 3x3 convs, and replace them (e.g. use setattr) with Canvas.Placeholder. Then call canvas.sample (for detailed arguments, you can follow the docs of the func) and canvas.replace.