FluxML / FluxML/FastAI.jl

Model registry

Open
#246 0 comments 1 reaction 0 assignees View on GitHub
api-proposal enhancement plans
Dominant language
Julia
Stars
592
Forks
49
PR merge metrics
No merged PRs in 30d

Description

In addition to existing feature registries, FastAI.jl will be getting a model registry.

The model registry should allow

- domain and third-party packages to add additional models when loaded;
- using models from different "backend" deep learning libraries, i.e.
Flux.jl, Lux.jl, PyTorch, Jax; and
- searching for models suitable for a specific task, fitting a computational budget,
or built in a specific deep learning library

## Examples

Some code examples to show how the model registry can be used:

### Loading models

Load a pretrained ResNet implemented in Metalhead.jl for transfer learning:

```julia
load(models()["metalhead/resnet18/head"], pretrained=true)
```

Load the ResNet as an untrained backbone for a different task

```julia
load(models()["metalhead/resnet18/backbone"], pretrained=false)
```

### Searching for models

Find models that take in preprocessed images:

```julia
filter(models(), input=ImageTensor{2})
```

Or find a suitable model for a supervised learning task directly:

```julia
task = SupervisedTask(_)/ImageSegmentation(_)/TabularClassificationSingle(_)
filter(models(), input=task.blocks.x, output=task.blocks.y)
```

List models implemented in PyTorch:

```julia
filter(models(), backend=:pytorch)
```

Find models of a certain size:

```julia
filter(models(), input=ImageTensor{2}, nparams=<(1000000))
```

### Training workflow

Since models in the registry are associated with block information, we can use them
to automatically construct task-specific models using the `taskmodel` API (possibly
extended by an additional `backend` argument).

```julia
config = models()["torchvision/resnet18/backbone"]
backbone = load(config)

task = ImageSegmentation(_)
# build the task-specific model
model = taskmodel(task, # includes info about required input and target block for task
config.backend # dispatch on the DL library used, here :pytorch
backbone,
config.input, # backbone input block: `ImageTensor{2}(3)`
config.output) # backbone output block: `ConvFeatures{2}(512)`

learner = tasklearner(task, data; model)
fit!(learner, 10)
```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.