Hello to everyone!
I have a base model which is a cycle GAN, I am using some layer of the generator to reuse learned features to fed a classifier. Train the two networks separately is not a problem. My doubt is related to how to build the cost function in order to optimize my cycleGAN based on how is performing the classifier.
Do you think this pseudo-code could be the right approach?
Both the loss functions optimize the base model w.r.t. the classification result. Which one is better?
basemodel = Mybasemodel()
basemodel.load_state_dict(torch.load('...'))
basemodel.train()
calssifier = Classifier()
calssifier.train()
loss_basemodel_fn = torch.nn...()
loss_classifier_fn = torch.nn...()
opt_basemodel = torch.optim...(basemodel.parameters())
# create an optimizer which optimizes both the networks
param_chain = itertools.chain(calssifier.parameters(), basemodel.parameters())
opt_classifier = torch.optim...(param_chain)
for i, data in enumerate(dataset):
# Load data
input = data['input']
label_basemodel = data['label_basemodel']
label_classifier = data['label_classifier']
# use part of the base model to get a feature map
feature_map = basemodel.model[first_layer:N_layer](input)
# get the prediction
pred = calssifier(feature_map)
output_base_model = basemodel.model(input)
opt_basemodel.zero_grad()
opt_classifier.zero_grad()
# BOTH THESE LOSSES WILL OPTIMIZE MY BASENODEL? WHICH ONE IS BETTER?
loss_classifier = loss_classifier_fn(pred, label_classifier)
# combine the two loss function so that the base model optimizer
# take into consideration the error of the calssifier when optimizes
# the basemodel
loss_basemodel = loss_basemodel_fn(output_base_model, label_basemodel) + loss_classifier
loss_basemodel.backward()
loss_classifier.backward()
opt_basemodel.step()
opt_classifier.step()