Added YoloV8 to examples

[Benny-Nottonson]

WIP

[andresnowak]

you can the submodules like how they did here (that would be the best I think, but i can try it later) yolov8 architecutre

[Benny-Nottonson]

you can the submodules like how they did here (that would be the best I think, but i can try it later) yolov8 architecutre

I will probably try this next, just need to figure out what is missing

[Benny-Nottonson]

you can the submodules like how they did here (that would be the best I think, but i can try it later) yolov8 architecutre

yeah this looks perfect now that I understand a little better, thank you

[andresnowak]

I have this here if you want

import basalt.nn as nn
from basalt import Tensor, TensorShape
from basalt import Graph, Symbol, OP, dtype



fn CSM(inout g: Graph, x: Symbol, out_channels: Int, kernel_size: Int, padding: Int, stride: Int, dilation: Int) -> Symbol:
    var conv = nn.Conv2d(g, x, out_channels, kernel_size, padding, stride, dilation)
    var sigmoid = g.op(OP.SIGMOID, conv)
    return g.op(OP.MUL, conv, sigmoid) # silu


fn C2f(inout g: Graph, x: Symbol, out_channels: Int, kernel_size: Int, padding: Int, stride: Int, dilation: Int) -> Symbol:
    var conv = nn.Conv2d(g, x, out_channels, kernel_size, padding, stride, dilation)
    var split = g.split(conv, 2, 1)

    @parameter
    fn bottleneck(x: Symbol, out_channels: Int, shortcut: Bool = False) -> Symbol:
        var conv1 = nn.Conv2d(g, x, out_channels, 1, 0, 1, 1)
        var conv2 = nn.Conv2d(g, conv1, out_channels, 3, 1, 1, 1)

        if shortcut:
            return g.op(OP.ADD, x, conv2)
        else:
            return conv2

    var y1 = bottleneck(split[0], out_channels, True)
    var y2 = bottleneck(split[1], out_channels, True)

    var y = g.concat(y1, y2, dim = 1)

    return CSM(g, y, out_channels, 1, 0, 1, 1)


fn SPPF(inout g: Graph, x: Symbol, out_channels: Int, kernel_size: Int, padding: Int, stride: Int, dilation: Int) -> Symbol:
    var conv = nn.Conv2d(g, x, out_channels, kernel_size, padding, stride, dilation)
    
    var maxpool2d_1 = nn.MaxPool2d(g, x, kernel_size=5, padding=2)
    var maxpool2d_2 = nn.MaxPool2d(g, x, kernel_size=5, padding=2)
    var maxpool2d_3 = nn.MaxPool2d(g, x, kernel_size=5, padding=2)

    var y = g.concat(conv, maxpool2d_1, maxpool2d_2, maxpool2d_3, dim = 1)

    return CSM(g, y, out_channels, 1, 0, 1, 1)


fn YoloV8(batch_size: Int) -> Graph:
    var g = Graph()
    var x = g.input(TensorShape(batch_size, 3, 640, 640))
    
    # Backbone
    var conv_1 = CSM(g, x, 32, 3, 1, 1, 1)
    var conv_2 = CSM(g, conv_1, 64, 3, 1, 1, 1)
    var C2f_1 = C2f(g, conv_2, 128, 3, 1, 1, 1)
    var conv_3 = CSM(g, C2f_1, 128, 3, 1, 1, 1)
    var C2f_2 = C2f(g, conv_3, 256, 3, 1, 1, 1)
    
    var conv_4 = CSM(g, C2f_2, 256, 3, 1, 1, 1)
    var C2f_3 = C2f(g, conv_4, 512, 3, 1, 1, 1)

    var conv_5 = CSM(g, C2f_3, 512, 3, 1, 1, 1)
    var C2f_4 = C2f(g, conv_5, 1024, 3, 1, 1, 1)
    var SPPF_1 = SPPF(g, C2f_4, 1024, 3, 1, 1, 1)


    # NOTE: This is where the "Resize: 4" OP is
 
    return g ^

But I think the sizes for the csm's are wrong (the stride, outchanels, etc..)

[andresnowak]

I put the function separately instead of using @parameter, for me it is easier to read in this way but you can change it if you want, but I think it is better this way.

[Benny-Nottonson]

I put the function separately instead of using @parameter, for me it is easier to read in this way but you can change it if you want, but I think it is better this way.

No you are right, I prefer that now that I see it. Im gonna change some variable names and restructure a little but nice changes