shivance

class RotaryEmbedding(keras.layers.Layer):

    def __init__(self, hidden_dim):
        super().__init__()
        self.hidden_dim = hidden_dim

    def build(self, input_shape):
        super().build(input_shape)
        self.inverse_freq = self.add_weight(
            "inverse_freq", shape=(self.hidden_dim // 2,), dtype=tf.float32

shivance

shivance

shivance

shivance

function (u::UNet)(x::AbstractArray)
    enc_out = []

    out = x
    for i in 1:4
        out = u.encoder[i](out)
        push!(enc_out, out)

        out = u.pool[i](out)
    end

shivance

def forward(self, x):
    enc1 = self.encoder1(x)
    enc2 = self.encoder2(self.pool1(enc1))
    enc3 = self.encoder3(self.pool2(enc2))
    enc4 = self.encoder4(self.pool3(enc3))

    bottleneck = self.bottleneck(self.pool4(enc4))

    dec4 = self.upconv4(bottleneck)
    dec4 = torch.cat((dec4, enc4), dim=1)

shivance

function upconv_block(in_chs::Int, out_chs::Int, kernel = (2, 2))
    return ConvTranspose(kernel, in_chs => out_chs; stride = (2, 2), init = _random_normal)
end

function UNet(in_channels::Integer = 3, inplanes::Integer = 32,
              outplanes::Integer = inplanes)
    features = inplanes

    encoder_layers = []
    append!(encoder_layers, [unet_block(in_channels, features)])

shivance

struct UNet
    encoder::Any
    decoder::Any
    upconv::Any
    pool::Any
    bottleneck::Any
    final_conv::Any
end
@functor UNet

shivance

function unet_block(in_chs::Int, out_chs::Int, kernel = (3, 3))
    return Chain(
             conv1 = Conv(kernel, in_chs => out_chs; pad = (1, 1),
                          init = _random_normal),
             norm1 = BatchNorm(out_chs, relu),
             conv2 = Conv(kernel, out_chs => out_chs; pad = (1, 1),
                          init = _random_normal),
             norm2 = BatchNorm(out_chs, relu)
          )
end

shivance

@staticmethod
def _block(in_channels, features, name):
    return nn.Sequential(
        OrderedDict(
            [
                (
                    name + "conv1",
                    nn.Conv2d(
                        in_channels=in_channels,
                        out_channels=features,