I have a UNet-like architecture that I want to parametrize (in terms of numbers of channels, kernel sizes, etc) and adapt for also taking an embedding vector as an input. While trying to code it I realized that there is too much indexing calculations, and coding it directly by creating every layer with right indicies and hyperparameters is tedious. Is there a better way?
I am talking about something like a graph that I could create in code by creating:
- verticies associated with intermidiate CxHxW tensors
- edges associated with convolutions, poolings, upconvolutions,skip-connections
It would be much easier to code, since its straightforward to write down what kind of intermidiate tensors will be there, but not as easy to write down creation of layers with correct hyperparams that connect correct layers, and also apply them in forward(). I want some general procedure to do that for me, so that I only have to specify general structure of the computational graph (i.e. which tensors are connected with which type of operation, and what shape they have).
Is there any general solution like that?