[Python] How to nest representer_mappings in ruamel.yaml?

EDIT: made the question more concrete

I want to have the concept of a learning task for optimizer benchmarking

@dataclass
class LearningTask
model: torch.nn.Module
loss: torch.nn.Module
data: lightning.LightningDataModule


Of course it might be nice to serialize this into yaml to have benchmarks be configured by yamls or to store the yaml configuration with the performance benchmark.

But the LearningTask above is ill suited for this purpose for two reasons:

1. I might want to repeat the Task with different seeds but the initialization of the parameter weights of the model happens at initialization
2. The random parameters of the model should not end up in the yaml configuration. But separating the generated and default parameters from the parameters passed to the constructor is very annoying after the object is constructed.

So instead, the LearningTask is not defined as the fully assembled object, but instead as a Factory:

from attr import dataclass
from typing import Any

@dataclass
class Params:
model: dict[str,Any]
loss: dict[str,Any]
data: dict[str,Any]

@dataclass
class BetterLearningTask:
model_class: type[torch.nn.Module]
loss_class: type[torch.nn.Module]
data_class: type[lightning.LightningDataModule]
params: Params

@property
model(self):
return self.model_class(**self.params.model)

# similar for loss and data

factory = BetterLearningTask(
model_class = MyModelInheritingFromPytorchModule,
loss_class = torch.nn.Crossentropy,
data_class = MyLightningDataModule
Params(model={"conv_kernel_size": 3}, loss={}, data={})
)


This formulation of the LearningTask is a Factory - it holds all the components to create the model, loss and data but not their initialized version. To store it, we only need to store the class names and the initialization parameters.

I want the yaml configuration of the factory to pretend like it is the simpler LearningTask, i.e.

!BetterLearningTask
model: !MyModelInheritingFromPytorchModule
conv_kernel_size: 3
loss: !torch.nn.Crossentropy {}
data: !MyLightningDataModule {}


What I have tried is the following

@yaml.register_class
@dataclass
class BetterLearningTask:
# ...

@classmethod
def to_yaml(cls, representer, node):
return representer.representer_mapping(
"!" + cls.__name__,
{
"model": representer.representer_mapping(
"!" + node.model_class.__name__,
node.params.model,
),
"loss": representer.representer_mapping(
"!" + node.loss_class.__name__,
node.params.loss,
),
}
)


but this results in

ruamel.yaml.representer.RepresenterError:
cannot represent an object: MappingNode(tag='!MyModelInheritingFromPytorchModule', value=[(conv_kernel_size, 3)])


I actually do not know what the different representers are supposed to do, which is why I am struggling to figure out what I should do differently. The documentation does not include the representers but the maintainers seems to be quite active answering questions here, so I thought I would try.

In this answer: np.array in yaml nesting of representer at least seem to work.

Continue reading...