Data block API

[lorenzoh]

It would be nice to have an API for easily constructing learning methods as manually implementing all the methods can get tedious and the resulting methods don't compose well.

The API would be similar to fastai's data block API with the main difference that it is limited to learning methods, i.e. it keeps data container creation and task-specific data encoding separate (handling only the encoding).

Based on Blocks which represent a kind of data and Encodings, transformations that encode data and are optionally invertible allowing the decoding of outputs.

API

Best to give an example of what using it would look like. Below are reimplementations of some of FastAI.jl's computer vision methods.

ImageClassificationSingle(sz, classes) = Method(
    blocks=(Image{2}(), Label(classes)),
    encodings=[
        ProjectiveTransforms(sz),
        ImagePreprocessing(),
        OneHot()
    ]
)

ImageSegmentation(sz, classes) = Method(
    blocks=(Image{2}(), Mask{2}(classes)),
    encodings=[
        ProjectiveTransforms(sz),
        ImagePreprocessing(),
        OneHot()
    ]
)

SiameseSimilarity() = Method(
    blocks=((Image{2}(), Image{2}()), Label([true, false])),
    encodings=[
        ProjectiveTransforms(sz),
        ImagePreprocessing(),
    ],
)

# TableClassification
tableblock = TableRow(table=traindf; catcols, contcols)  # construct block with vocabulary from DataFrame
Method(
    blocks=(tableblock, Label(classes)),
    encodings=[
        TableTransforms(),
    OneHot()
    ]
)

Given just these short definitions and the block and encoding definitions and the right interfaces inplace, it would be possible to derive the following:

• core interface (encoding, decoding, incl. buffered versions)
• validation of input data
• plotting interface
• model building based on input and target blocks
• loss functions based on target block

By grouping functionality by block or encoding, it would be much easier to compose and reuse different steps.

Status

Implemented:

• interfaces
  • Block
  • Encoding
  • StatefulEncoding
• blocks
  • Image
  • Mask
  • Label
  • LabelMulti
  • OneHotTensor
  • ImageTensor
• encodings
  • ProjectiveTransforms (now works with 3D images, masks and keypoints)
  • ImagePreprocessing (now works with 3D images)
  • OneHot (for labels, multi-class labels and masks)

To-do:

• interfaces:
  • splitting Encoding into AbstractEncoding, Encoding and WrapperEncoding
  • splitting Block into AbstractBlock, Block and WrapperBlock
  • data block learning method
  • plotting interface
  • model construction interface
• encodings:
  • Only
  • TaggedBlock
• replace old learning method definitions with new API

The encodings depend on some minor changes to DataAugmentation.jl, to be released soon.

How do I

Apply an encoding to multiple blocks

By default, an encoding transforms every block for which an encode method is implemented. For example encode(ProjectiveTransforms(...), _, (Image(), Mask()), (img, mask) will encode both image and mask (with the same random state for the augmentations) while encode(ProjectiveTransforms(...), _, (Image(), Label(classes)), (img, class) will encode only the image, and the class is passed through unchanged, since no method encode(::ProjectiveTransforms, _, ::Label, _) is implemented.

Apply an encoding to a specific block only

Let's say you have blocks of different types and an encoding implemented for all blocks, but you only want to encode a single block. This could be achieved with a wrapper encoding that only applies the wrapped encoding if a condition is met. The below example shows how ProjectiveTransforms, which would encode Image and Mask is wrapped so only Images are transformed.

Method(
    blocks=(Image{2}(), Image{2}()),
    encodings=[
        Only(Image, ProjectiveTransforms()),
    ]
)

Now what if you had multiple blocks of the same type? We need a way to select which blocks to transform and which to leave be, but can no longer use the type to distinguish the blocks. Note that we cannot use indices of a tuple of blocks as selectors since the same set of encodings need to be callable on different sets of blocks (for example, during training inputs and targets are encoded together, and during inference inputs are encoded by themselves and model outputs are also decoded by themselves).

One solution is to introduce a wrapper block (yes, Julia is big on composition) that associates a tag with the block which can then be referenced in the encoding wrapper.

Method(
    blocks=(Tagged(:encodeme, Image{2}()), Image{2}()),
    encodings=[
        Only(:encodeme, ProjectiveTransforms()),
    ]
)

Write an encoding that combines multiple blocks

By default, applying an encoding to a tuple of blocks will apply the encoding to each block individually. This can be overwritten by implementing an encode method that dispatches on Tuple and can combine multiple blocks.

Below example shows a transform that concatenates selected blocks:

encoding = Concat(Image, dim=3)  # concats all image blocks, may also use a tag as selector, see above
# implements `encode(::Concat, _, blocks::Tuple, datas::Tuple) -> Image`

Method(
    blocks=(Image{2}(), Image{2}()),
    encodings=[encoding]
)

Apply multiple encodings to the same block

Can be done with a wrapper transform that stores multiple encodings.

Method(
    blocks=Image{2}(),
    encodings=[
        ProjectiveTransforms(...),   # returns `Image`
        Encodings(ImagePreprocessing(), identity)  # returns (`ImageTensor`, `Image`)
])

Create a learning method where model output block differs from the encoded target block

This has come up for me during segmentation where I used a custom loss function to weigh foreground losses, So instead of the loss function being loss(y_pred, y) it was weightedloss(y_pred, (y, weights)). Let's say we have an encoding CreateForegroundWeights that transforms a Mask block to create a Weights block. We can use the above Encodings to apply one-hot encoding and weight creation together (Encodings(OneHot(), CreateForegroundWeights())). If we transform the blocks (Image{2}(), Mask{2}()) the output would be xblock, yblock = (ImageTensor(), (Mask(), Weights()). We can see that our ys would be compatible with the loss function, great. However, model outputs will be just Mask blocks which leads to a problem when decoding since by default the method expects the same target block (yblock). We can use the outputblock keyword argument to override this.

ImageSegmentation(sz, classes) = Method(
    blocks=(Image{2}(), Mask{2}(classes)),
    encodings=[
        ProjectiveTransforms(sz),
        ImagePreprocessing(),
        Encodings(OneHot(), CreateForegroundWeights()),
    ],
    outputblock = Mask{2}()  # defaults to `encodedblock(encodings, blocks[2]) = (Mask{2}(), Weights{2}())`
)

Use blocks from different applications together

The API is application-agnostic so e.g. computer vision and tabular blocks can be used together. There is no special logic for vision methods. There are no default encodings associated with blocks as in fast.ai, every encoding is explicit.

[AriMKatz]

I crossposted your request for feedback to the FastAI discord and received the following response:

A couple thoughts:

From the design spec, it looks encodings cannot be applied separately to each input, and the API is limited to one set of encodings. For self supervised learning and other multi-image tasks, it's often necessary to apply different sets of augmentations to each image pair: one weakly augmented and the other strongly augmented. From the start I would design FastAI.jl's datablock to support this.

On a similar line of thought, in fastai2's datablock mixing input types, images + tabular for example, is not a straightforward task and requires creating separate image and tabular datablocks then mixing their dataloaders. Instead of adding an image and tabular block in one datablock. Would be nice if FastAI.jl's datablock could support this too.

[lorenzoh]

The comment is now adressed above under "Apply an encoding to a specific block only"