Can these three classes (stem, identity_classifier and colour_classifier) be integrated in a network? Thank you! For example:

[junzai0215]

Can these three classes (stem, identity_classifier and colour_classifier) be integrated in a network? Thank you! For example:

import numpy as np
from itertools import chain
import torch.nn
import pytorch_revgrad

class Classifiers(nn.Module):
    def __init__(self):
        super(Classifiers, self).__init__()
        self.stem = torch.nn.Sequential(
           torch.nn.Linear(128, 256),
           torch.nn.ReLU(),
           torch.nn.Linear(256, 512),
           torch.nn.ReLU(),
           torch.nn.Linear(512, 128),
           torch.nn.ReLU(),
           torch.nn.Linear(128, 64),
        )

       self.identity_classifier = torch.nn.Sequential(
          torch.nn.Linear(64, 64),
          torch.nn.ReLU(),
          torch.nn.Linear(64, 10),
       )

       self.colour_classifier = torch.nn.Sequential(
          pytorch_revgrad.RevGrad(),
          torch.nn.Linear(64, 64),
          torch.nn.ReLU(),
          torch.nn.Linear(64, 2),
       )
      def forward(self, inp):
         intermediate_features = self.stem(inp)
         identity_logits = self.identity_classifier(intermediate_features)
         colour_logits = self.colour_classifier(intermediate_features)
         return identity_logits, colour_logits

for epoch in range(100):
    for inp, iden, col in loader:
        identity_logits, colour_logits  = Classifiers(inp)
        identity_loss = torch.nn.functional.cross_entropy(identity_logits, iden)
        colour_loss = torch.nn.functional.cross_entropy(colour_logits, col)
        total_loss = identity_loss + alpha * colour_loss
        total_loss.backward()
        ......
  

Originally posted by @junzai0215 in #5 (comment)