I have this architecture, made of 3 NNs: enter image description here In code:

class VGGBlock(nn.Module):
    def __init__(self, in_channels, out_channels,batch_norm=False):


        conv2_params = {'kernel_size': (3, 3),
                        'stride'     : (1, 1),
                        'padding'   : 1

        noop = lambda x : x

        self._batch_norm = batch_norm

        self.conv1 = nn.Conv2d(in_channels=in_channels,out_channels=out_channels , **conv2_params)
        self.bn1 = nn.BatchNorm2d(out_channels) if batch_norm else noop

        self.conv2 = nn.Conv2d(in_channels=out_channels,out_channels=out_channels, **conv2_params)
        self.bn2 = nn.BatchNorm2d(out_channels) if batch_norm else noop

        self.max_pooling = nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2))

    def batch_norm(self):
        return self._batch_norm

    def forward(self,x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = F.relu(x)

        x = self.conv2(x)
        x = self.bn2(x)
        x = F.relu(x)

        x = self.max_pooling(x)

        return x
class VGG16(nn.Module):

  def __init__(self, input_size, num_classes=1,batch_norm=False):
    super(VGG16, self).__init__()

    self.in_channels,self.in_width,self.in_height = input_size

    self.block_1 = VGGBlock(self.in_channels,64,batch_norm=batch_norm)
    self.block_2 = VGGBlock(64, 128,batch_norm=batch_norm)
    self.block_3 = VGGBlock(128, 256,batch_norm=batch_norm)
    self.block_4 = VGGBlock(256,512,batch_norm=batch_norm)

  def input_size(self):
      return self.in_channels,self.in_width,self.in_height

  def forward(self, x):

    x = self.block_1(x)
    x = self.block_2(x)
    x = self.block_3(x)
    x = self.block_4(x)
    x = torch.flatten(x,1)

    return x
class VGG16Classifier(nn.Module):

  def __init__(self, num_classes=1,classifier = None,batch_norm=False):
    super(VGG16Classifier, self).__init__()

    self._vgg_a = VGG16((1,32,32),batch_norm=True)
    self._vgg_b = VGG16((1,32,32),batch_norm=True)
    self._vgg_star = VGG16((1,32,32),batch_norm=True)
    self.classifier = classifier

    if (self.classifier is None):
        self.classifier = nn.Sequential(
          nn.Linear(2048, 2048),
          nn.Linear(2048, 512),
          nn.Linear(512, num_classes)

  def forward(self, x1,x2,x3):
      op1 = self._vgg_a(x1)
      op2 = self._vgg_b(x2)
      op3 = self._vgg_star(x3) 
      x1 = self.classifier(op1)
      x2 = self.classifier(op2)
      x3 = self.classifier(op3)

      return x1,x2,x3

      return xc
model1 = VGG16((1,32,32),batch_norm=True)
model2 = VGG16((1,32,32),batch_norm=True)
model_star = VGG16((1,32,32),batch_norm=True)
model_combo = VGG16Classifier(model1,model2,model_star)

I have to implement this custom loss: enter image description here In code:

class CombinedLoss(nn.Module):
    def __init__(self, loss_a, loss_b, loss_star, _lambda=1.0):
        self.loss_a = loss_a
        self.loss_b = loss_b
        self.loss_star = loss_star


    def forward(self,y_hat,y):

        return (self.loss_a(y_hat[0],y[0]) + 
                self.loss_b(y_hat[1],y[1]) + 
                self.loss_combo(y_hat[2],y[2]) + 
                self._lambda * torch.sum(model_star.weight - torch.pow(torch.cdist(model1.weight+model2.weight), 2)))

Maybe is not well written this loss (i accept advices :D) Now, my question is: considering the fact that i have a model made of 3 networks, and a loss which have 3 labels (y[0, y1 and y2) I have to split my dataset in 3 parts? For example, if I use MNIST, I have to split MNIST in MNIST_A, MNIST_B and MNIST_STAR, creating a loader for each of these 3 datasets, and pass to the model when I have to train?

  • $\begingroup$ Yes, you can create a single data loader that will generate 3 inputs for your model. $\endgroup$ Jun 15 at 16:39
  • $\begingroup$ You say, don't split the dataset in 3 parts, but create a dataloaders which pass 3 inputs? $\endgroup$
    – CasellaJr
    Jun 15 at 16:48
  • $\begingroup$ Actually, I cannot understand what do you mean by splitting the dataset. If you can show how your data looks then I can actually write code for the data loader. $\endgroup$ Jun 15 at 17:02
  • $\begingroup$ MNIST, my dataset is MNIST. I mean how can I write a training function in which I pass data to each model? $\endgroup$
    – CasellaJr
    Jun 15 at 18:25
  • $\begingroup$ Because basically, i have model_combo that contains model_a, model_b and model_star, and I want the first 33% of dataset (training, val and test) for model_a, then the other 33% for model_b and last 33% for model_star. So, I want that each model has a different Subset of MNIST $\endgroup$
    – CasellaJr
    Jun 15 at 18:26

No need to split the dataset. You can do this on runtime and using a single data loader. Following is the code for the same and explanations included in comments -

# imports
import torch
from torch.utils.data import Dataset
from torchvision import datasets
from torch.utils.data import DataLoader

# This is how I am downloading MNIST dataset, I use the .pt files from the downloaded dataset.
training_data = datasets.MNIST(

# A single custom dataset
class MnistImageDataset(Dataset):
    def __init__(self):
        # first the dataset will be loaded
        mnist_train = torch.load("data/MNIST/processed/training.pt")

        self.m_images = mnist_train[0] # all images
        self.m_labels = mnist_train[1] # all labels

        self.no_labels = len(self.m_labels)//3 # len of images a single model will be trained with

    def __len__(self):
        # dataset length will be the total_len/3
        return len(self.m_labels)//3

    def __getitem__(self, idx):
        # create dictionaries
        images = {}
        labels = {}

        # select images from each split for each model
        images['model_a'] = self.m_images[idx]
        images['model_b'] = self.m_images[idx+self.no_labels]
        images['model_star'] = self.m_images[idx+self.no_labels*2]

        # select labels from each split for each model
        labels['model_a'] = self.m_labels[idx]
        labels['model_b'] = self.m_labels[idx+self.no_labels]
        labels['model_star'] = self.m_labels[idx+self.no_labels*2]

        return images, labels

# create dataset and wrap into dataloader
mnist = MnistImageDataset()
train_dataloader = DataLoader(mnist, batch_size=64)

for images, labels in train_dataloader:
        # sample inference
        op1,op2,op3 = model_combo(torch.unsqueeze(images['model_a'],1),torch.unsqueeze(images['model_b'],1),torch.unsqueeze(images['model_star'],1))
        # next comes your loss and other pipeline
  • $\begingroup$ Thanks for the answer. However, I think I will procede in another way for the moment $\endgroup$
    – CasellaJr
    Jun 16 at 15:10

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.