How to use Metalhead'model as ImageClassifier model?

[zsz00]

How to use Metalhead'model as ImageClassifier model?
such as:

using Flux, MLJFlux, Metalhead

ImageClassifier = @load ImageClassifier

model = ResNet50(pretrain=false, nclasses=20)

clf = ImageClassifier(builder=model,
		                  epochs=500,
                                  optimiser=ADAM(0.001),
		                  loss=Flux.crossentropy,
                                  batch_size=512,
                                  acceleration=CUDALibs(),)

mach = machine(clf, images, labels)

[zsz00]

I found it

mutable struct MyConvBuilder
end

function MLJFlux.build(b::MyConvBuilder, rng, n_in, n_out, n_channels)
    model = ResNet18(pretrain=false, nclasses=n_out)
    return model
end

clf = ImageClassifier(builder=MyConvBuilder(),
		                  epochs=500,
                                  optimiser=ADAM(0.001),
		                  loss=Flux.crossentropy,
                                  batch_size=512,
                                  acceleration=CUDALibs(),)

[ablaom]

Thanks. Perhaps we can now resolve #162! I'm happy to make MetalHead as a dependency of MLJFlux. PR welcome.

[ablaom]

Closed in favour of #205