Hi, I have a weird question about some crazy idea I want to materialize using PyTorch.
The problem: I want to train re-id model using multiple datasets. Such a solution already exists in https://github.com/KaiyangZhou/deep-person-reid 2 and it works.
However, I beleive it is not optimal, as in this repo they combine the datasets, sums all unique entities and relabel all of them. Example: I have 3 datasets with 50,000 unique entities each. When using
deep-person-reid I will get 150,000 unique entities at the end and the last softmax layer will have a size of 150,000 neruons.The problem is that such approach causes explosion of the dimension of the softmax layer.
In my case I have almost 1 million unique entities, which casue huge GPU RAM overhead just to store the model! (no worries, my entities are not human people  )
To alleviate the problem I came up with a crazy (I think it is a little bit crazy).
My idea:
(For all examples let’s assume I have 3 different datasets, with 50,000 entities and image sizes etc. are consistent across all datasets)

Create separate dataloaders, which will be used in turns during training. (3 seperate dataloaders)
Create removable softmax layer that would be used in conjuntion with a specific dataset/dataloader. (3 seperate softmax layers)
Create separate (?) optimizers that would be used in conjuction with a specific softmax layer. (3 separate optimizers)

In such a setup the training would look like this:
0. Move model (without softmax layer), criterion etc. to GPU.

Get 1st dataloder & optimizer.
Append 1st softmax layer to the model.
Run batch
Update parameters.
Remove 1st softmax layer from model and from GPU
Get 2nd dataloader & optimizer
Append 2nd softmax layer to the model.
Run batch
Update parameters.
Remove 2nd softmax … and so on.

I see here 2 main challenges:


I could not find a clear explanation how optimazers works under-the-hood, but I heard that they are somehow binded to all modules in the model.
I assume there should be also 3 seperate optimizers that would be dedicated to each softmax layer. Moreover, each optimizer should be disabled during not-their-training cycle
(so that not accumulate gradient when their dedicated softmax layer is not used in current training cycle).
Maybe write a custom optimizer that would allow to freeze its parameters when needed?
What would be a sensible way to approch this challenge?


Not sure how to deal with the problem of replacing softmax layer during training.
My instant idea was to just choose desired softmax layer in forward function, but then other softmax layers would be still laying on the GPU RAM,
which would be the same effectively as implemented in deep-person-reid, thus it would not solve the main problem of consuming too much GPU RAM.
Second thought was to rather modify the model architecture - replace the softmax layer only (with corresponding weights ofc), however I am not
sure how doble it is on-the-fly.
Moreover there will be some performance overhead with this approach probably.


Hope, I explained the idea clearly enough. I am more than happy to explain it more if needed 
Most of all I hope someone can direct me to docs/tutorial/something else that would give me more insight in the inner working of PyTorch, which help me to solve my problem.