search

pxl-th / Segmentation.jl

Segmentation models in Julia
14 stars 2 forks source link

Compatibility with Metalhead #1

Open DhairyaLGandhi opened 2 years ago

DhairyaLGandhi commented 2 years ago

It should be possible to maintain compatibility with most common models used in CV, so what do you think of having examples using segmentation with models directly from Metalhead. Also, it would be good to move EfficientNet to Metalhead. That way we could use it to complete some model classes, add pretrained weights and include it in benchmarking in FluxBench.jl.

pxl-th commented 2 years ago

It should be possible to maintain compatibility with most common models used in CV, so what do you think of having examples using segmentation with models directly from Metalhead.

Integrating with Metalhead would be a good idea. The only thing is where to implement stages extraction for the models. Right now both ResNet and EfficientNet implement it in their packages.

Also, it would be good to move EfficientNet to Metalhead.

Agree, especially since I've confirmed that it can train :) Also it would be good to get rid of PyTorch's weights conversion stuff. I can create PR once I have some free time.

DhairyaLGandhi commented 2 years ago

We'll need a decent way of ensuring we can share weights going from pytorch to flux and the reverse direction too, so we might retain that for the time being. I'd appreciate some help updating Torch.jl to use the latest version! That would help quite a bit!

I'm not sure what the stages conversion is doing? It seems like it should be possible to simplify which layers are to be run. On a separate note, head and top should be identity instead of nothing by default. That would help generalise the single argument forward pass.