Hello everyone, hope you are having a great day.
How can I have layer wise training in Pytorch? I mean, suppose I have a network that trains like normal
but parts of the network also gets optimized independently ? its something like this
|________|
|module1 |
|________|
|module2 |
|________|
|module3 |
|________|
It goes like this I feed the model like normal :
for imgs, labels in dataloader:
imgs = imgs.to(device)
labels = labels.to(device)
preds = model(imgs)
loss = criterion(preds, labels)
....
but what is not shown here is that each module can have its own criterion,
each module receives a batch of images for example, runs the normal thing it does
and it also trains until the loss reaches to a specific quantity then gets out of the loop for example and
sends its outputs to the next layer, this goes on until the network predictions are made available which
then the loss is calculated and the network weights is updated.
Thank you all very much in advance