Help to extend the model in Metalhead.jl

yeruo
Sat, 12 Aug 2023 10:04:19 +0000
General Usage

I implemented a self-supervised learning algorithm, MAE (Masked Auto Encoder), based on Flux.jl and Metalhead.jl, It seems to work, especially when propagating forward, but during gradient descent, Zygote alerted me that the array had been mutated in my code.

My code is as follows:

the Decoder and contruction

struct MaskDecoder
    projection
    mask_token
    decoder
end

@functor MaskDecoder
"""
    MaskDecoder(
        imsize::NTuple{N,Int64,},
        patch_size::NTuple{N,Int64},
        encodplanes::Integer,
        pretrain::Bool=false,
        outchannels::Integer=3;
        embedplanes::Integer=512,
        depth::Integer=8, nheads::Integer=16
    )

 A decoder for Masked Auto Encoder

# Arguments

- `imsize`: the input image size.
- `patch_size`: the size of the patch.
- `encodplanes`: the number of channels of ViT encoder's output.
- `pretrain`: whether to load pretrained paramters.
- `outchannels`: the number of channels in the recovered image.
- `embedplanes`: the number of channels projected to decoder transformer.
- `depth`: the number of blocks in decoder transformer.
- `nheads`: number of attention heads in transformer.
"""
function MaskDecoder(
    imsize::NTuple{N,Int64,},
    patch_size::NTuple{N,Int64},
    encodplanes::Integer,
    pretrain::Bool=false,
    outchannels::Integer=3;
    embedplanes::Integer=512,
    depth::Integer=8, nheads::Integer=16
) where {N}
    npatches = prod(imsize .÷ patch_size) + 1
    model = MaskDecoder(
        Dense(encodplanes, embedplanes),
        zeros32(embedplanes, 1, 1), # the <MSK>
        Chain(
            ViPosEmbedding(embedplanes, npatches),
            transformer_encoder(embedplanes, depth, nheads),
            Dense(embedplanes, prod(patch_size) * outchannels)
        )
    )
    if pretrain
        throw(ErrorException("MAE Decoder do not have the pretrain paramters."))
    end
    return model
end

6 posts - 1 participant