Representing a neural network as a resource JSON for API?

Question

I want to represent a neural network as a resource available to a REST API. Whenever a client wants to create a neural network model, they can POST a JSON representation of the neural network, which should be a Model object that contains a list of Layer objects.

A Model's primary data is that list of Layer objects, where each Layer can be a fully-connected Linear layer, some sort of activation layer like ReLU, or some other valid PyTorch network layer like Conv2d. However, each of these layers requires a different set of parameters. The Linear layer would require the number of nodes, Conv2d requires kernel size and channel count, while ReLU wouldn't really require anything.

My question is, how should I go about representing this list of Layers, specifically in FastAPI?

I had two ideas, but I'm hesitant about both.

Idea 1: Enum parameter with params variable

Here, we have a single Layer model that represents everything, with "types" handled by the type variable.

from enum import Enum
from pydantic import BaseModel

class LayerType(str, Enum):
    linear = "Linear"
    relu = "ReLU"
    conv2d = "Conv2d"
    # ... other layer types

class Layer(BaseModel):
    type: LayerType
    params: dict # ?

class Model(BaseModel): # main representation of the neural network
    name: str
    layers: list[Layer]

With this idea, the client would specify a list of JSON objects, specifying type and filling in the correct information in params. However, this would make type-checking extremely tedious since I would have to validate params myself and/or specify a long @model_validator() method, right?

Idea 2: Separate API Model for each Layer Type

Instead of that, I can specify a Model for each layer type like this:

from pydantic import BaseModel
from typing import Union
from typing_extensions import Self

class LinearLayer(BaseModel):
    size: int

    @model_validator(mode='after')
    def check_params(self) -> Self:
        assert self.size > 0 and self.size < 256, "`size` must be between 1 and 255, inclusive."
        return self

class DropoutLayer(BaseModel):
    dropout_prob: float = 0.5

    @model_validator(mode='after')
    def check_params(self) -> Self:
        assert self.dropout_prob >= 0.0 and self.dropout_prob <= 1.0, \
        "`dropout_prob` must be between 0.0 and 1.0, inclusive."
        return self

class Conv2DLayer(BaseModel):
    num_channels: int
    kernel: int | None = 3

    @model_validator(mode='after')
    def check_params(self) -> Self:
        assert self.num_channels > 0 and self.num_channels < 256, \
                "`num_channels` must be in the interval [1, 255]"
        assert self.kernel > 0 and self.kernel < 4, \
                "`kernel` must be in [1, 3]"
        return self

Layer = Union[LinearLayer, DropoutLayer, Conv2DLayer] # ***

class Model(BaseModel): # main representation of the neural network
    name: str
    layers: list[Layer]

Between these two ideas, this idea seems definitely a lot better since I can check each field and limit its range per layer type. One minor issue would be that the Layer type would need to be a very long Union if I wanted several different Layers. But the main issue I had with this idea was getting it to correctly match each layer. If I have a simple API POST method like:

from fastapi import FastAPI
from .schemas import * # the Model schema and layer stuff

app = FastAPI()
@app.post("/models/")
async def create_model(model: Model) -> Model:
    return model

And test this POST method with this simple JSON:

{
  "name": "test_model",
  "layers": [
    {
      "size": 0
    },
    {
      "dropout_prob": 0.5
    },
    {
      "num_channels": 0
    }
  ]
}

I get the following response:

{
  "name": "test_model",
  "layers": [
    {
      "dropout_prob": 0.5
    },
    {
      "dropout_prob": 0.5
    },
    {
      "dropout_prob": 0.5
    }
  ]
}

I think I'm not understanding the fundamental principles of how FastAPI matches the JSON object with the corresponding pydantic BaseModel. This is my first time working with REST APIs.

Is there some kind of pydantic inheritance/polymorphism thing I can use to solve this problem? Or should I go about implementing idea #1 instead of #2? I haven't had much luck looking through the FastAPI and pydantic documentation, or finding any mention of a situation like this.

Answer 1

One possible solution cloud be a combination of your two suggestions using Pydantic's discriminated union. Then you would do something along these lines:

from pydantic import BaseModel, Field
from typing import Annotated, Literal, Union


class LinearLayer(BaseModel):
    layer_type: Literal["Linear"]
    size: int

    ...

class DropoutLayer(BaseModel):
    layer_type: Literal["ReLU"]
    dropout_prob: float = 0.5

    ...

class Conv2DLayer(BaseModel):
    layer_type: Literal["Conv2d"]
    num_channels: int
    kernel: int | None = 3

    ...

Layer = Annotated[
    Union[LinearLayer, DropoutLayer, Conv2DLayer],
    Field(discriminator="layer_type")
]

class Model(BaseModel): # main representation of the neural network
    name: str
    layers: list[Layer]

Then Pydantic will use the layer_type field to infer which model class to use, and use that to parse the other parameters.