Generator Train

[way-sal]

Thanks for sharing your code, but I didn't understand where is the generator training step and Loss function.

[JinY97]

        p_t = model(x_t).view(x_t.shape[0], -1)

        fake_y, _, fake_feature2, _ = model_d(p_t.unsqueeze(dim=1))
        _, _, true_feature2, _ = model_d(y_t.unsqueeze(dim=1))

        y_t = y_original
        loss = denoise_loss_mse(p_t, y_t)
        
        if loss_type == "cls":
            g_loss = loss + w_c * (torch.mean((fake_y - 1) ** 2)) 
        elif loss_type == "feature": 
            g_loss = loss + w_f * denoise_loss_mse(fake_feature2, true_feature2)
        elif loss_type == "feature+cls": 
            g_loss = loss + w_f * denoise_loss_mse(fake_feature2, true_feature2) + w_c * (torch.mean((fake_y - 1) ** 2))  

        optimizer_D.zero_grad()
        optimizer.zero_grad()
        g_loss.backward()
        optimizer.step()
            
        losses.append(g_loss.detach())

“g_loss” is the generator loss.