Combining Multiple Models and Datasets

rsomani95: <p>I’m trying to combine multiple related datasets and models into one giant model. Here’s what the architecture looks like:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/7/47a797c252176f2c69c1ec690b94fca1458db816.png" data-download-href="https://discuss.pytorch.org/uploads/default/47a797c252176f2c69c1ec690b94fca1458db816" title="Screenshot 2020-05-23 at 11.58.54 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/7/47a797c252176f2c69c1ec690b94fca1458db816_2_345x221.png" alt="Screenshot 2020-05-23 at 11.58.54 PM" data-base62-sha1="adSXUBBtcXA9Yje7pZBb0ZHJlYi" width="345" height="221" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/7/47a797c252176f2c69c1ec690b94fca1458db816_2_345x221.png, https://discuss.pytorch.org/uploads/default/optimized/3X/4/7/47a797c252176f2c69c1ec690b94fca1458db816_2_517x331.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/4/7/47a797c252176f2c69c1ec690b94fca1458db816_2_690x442.png 2x" data-dominant-color="F2F2F2"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2020-05-23 at 11.58.54 PM</span><span class="informations">732×470 8.73 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>For each <code>head</code>, there’s a dataset with the following structure:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/7/97cf3e9e16f0ba2db84cb6241198892df37bdc4d.png" data-download-href="https://discuss.pytorch.org/uploads/default/97cf3e9e16f0ba2db84cb6241198892df37bdc4d" title="Screenshot 2020-05-24 at 12.01.52 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/7/97cf3e9e16f0ba2db84cb6241198892df37bdc4d_2_138x250.png" alt="Screenshot 2020-05-24 at 12.01.52 AM" data-base62-sha1="lEY7hQgeZs33ttiGMWlELGz7TwV" width="138" height="250" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/7/97cf3e9e16f0ba2db84cb6241198892df37bdc4d_2_138x250.png, https://discuss.pytorch.org/uploads/default/optimized/3X/9/7/97cf3e9e16f0ba2db84cb6241198892df37bdc4d_2_207x375.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/9/7/97cf3e9e16f0ba2db84cb6241198892df37bdc4d_2_276x500.png 2x" data-dominant-color="F0F0F0"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2020-05-24 at 12.01.52 AM</span><span class="informations">360×652 7.43 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I’ve referred to the following 2 sources for setting up the model, loss, and optimiser:</p><aside class="quote" data-post="27" data-topic="28383">
  <div class="title">
    <div class="quote-controls"></div>
    <img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/combining-trained-models-in-pytorch/28383/27">Combining Trained Models in PyTorch</a> 
  </div>
  <blockquote>
    The gradients will be calculated for the parameters of all submodules as well as the final classifier, if you didn’t disable it via .requires_grad=False. 
Also the state_dict() will return all parameters of all submodules and the final classifier.
  </blockquote>
</aside>
<aside class="quote" data-post="7" data-topic="45637">
  <div class="title">
    <div class="quote-controls"></div>
    <img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/muammar/48/76367_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/merging-two-models/45637/7">Merging two models</a> <a class="badge-category__wrapper " href="/c/vision/5"><span data-category-id="5" style="--category-badge-color: #AB9364; --category-badge-text-color: #FFFFFF;" data-drop-close="true" class="badge-category " title="Topics related to either pytorch/vision or vision research related topics"><span class="badge-category__name">vision</span></span></a>
  </div>
  <blockquote>
    <a class="mention" href="/u/alex.veuthey">@alex.veuthey</a>  I don’t. What I do is I forward propagate model A (an autoencoder), and subsequently use the latent hidden layer as input for model B (everything online?). I finally got time to keep working again on this and am to the point where I concatenate the parameters of the models in one single list.
  </blockquote>
</aside>
<p>
I’ve setup my model roughly as follows:</p>
<pre data-code-wrap="python"><code class="lang-python">class CombinedModel(nn.Module):
    def __init__(self):
        super(CombinedModel, self).__init__()
        self.features = models.mobilenet_v2(pretrained=True)
        self.pool   = nn.Sequential(*[
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(1)
        ])
        self.head1 = nn.Linear(1280,2)
        self.head2 = nn.Linear(1280,3)

    def forward(x, head_name):
        assert head_name in ['head1', 'head2']
        x = self.features(x)
        # only run the input for the head that is assigned
        # to that dataset
        return getattr(self, head_name)(x)
</code></pre>
<p>What I’d like to do is a round-robin style training where I train one batch from one dataset, fine-tune the respective head and the feature extractor’s parameters, then do another batch from another dataset, and so on.</p>
<p>I’ve setup individual loss functions for each head and a single optimiser for all of them. I’ve checked that the optimiser only calculates gradients for the <code>head_name</code> that is passed in while calling the <code>forward</code> method.</p>
<pre data-code-wrap="python"><code class="lang-python">optimiser = optim.Adam([
    {'params': model.features[-5:].parameters(), 'lr':1e-5},
    *[{'params':getattr(model,head).parameters(), 'lr':1e-3} 
      for head in ['head1', 'head2']
])

criterions = {
    'head1': nn.CrossEntropyLoss(),
    'head2': nn.CrossEntropyLoss()
}
</code></pre>
<p>I can create a dataset for the respective head with <code>torchvision.datasets.ImageFolder</code> but am unclear on how to combine it for multiple datasets since I’d like to switch between datasets after every batch.</p>
<p>I’d imagined the training loop to look something like this:</p>
<pre data-code-wrap="python"><code class="lang-python">model = CombinedModel()

# training mode
# each batch should switch (randomly) between
# head1 and head2 
for input,labels,head_name in dataloaders:
    # head_name == 'head1' or 'head2'
    optimiser.zero_grad()
    outputs = model(inputs, head_name)

    # calculate loss for respective head
    loss = criterions[head_name](outputs, labels)
    loss.backward()
    optimiser.step()
</code></pre>
<p>I’d love any input on what would be the right way to implement this.<br>
Thank you for your time!</p>

ptrblck: <p>One possible approach would be to create two separate <code>DataLoader</code> and use their <code>iterators</code> directly via:</p>
<pre><code class="lang-python">dataset_head1 = TensorDataset(torch.randn(10, 1), torch.randn(10, 1))
dataset_head2 = TensorDataset(torch.randn(10, 1), torch.randn(10, 1))

for epoch in range(2):
    print('epoch ', epoch)
    # in epoch loop
    loader1 = DataLoader(dataset_head1)
    iter_loader1 = iter(loader1)
    loader2 = DataLoader(dataset_head2)
    iter_loader2 = iter(loader2)
    
    try:
        while True:
            # train head1
            data, target = next(iter_loader1)
            print('training head1')
    
            # train head2
            data, target = next(iter_loader2)
            print('training head2')
    except StopIteration:
        pass
</code></pre>
<p>To avoid code duplication, you could of course write a <code>train</code> method and pass the current data into it, if that fits your use case.</p>
<p>Let me know, if this approach would work for you.</p>

rsomani95: <p>Excellent! This is definitely the right approach.</p>
<p>I suppose this is more of a python question, but with this exact approach, the training loop exits on the first <code>StopIteration</code>, which is a problem if the two dataset sizes differ. Is the only way to handle this in the except block?</p>
<p>Here’s a simulation of that situation from your code block:</p>
<pre><code class="lang-python">dataset_head1 = TensorDataset(torch.randn(10, 1), torch.randn(10, 1))
dataset_head2 = TensorDataset(torch.randn(2, 1), torch.randn(2, 1))

for epoch in range(2):
    print('epoch ', epoch)
    # in epoch loop
    loader1 = DataLoader(dataset_head1)
    loader2 = DataLoader(dataset_head2)
    
    iter_loader1 = iter(loader1)  # length - 10
    iter_loader2 = iter(loader2)  # length - 2
    
    try:
        while True:
            # train head1
            data, target = next(iter_loader1)
            print('training head1')
    
            # train head2
            data, target = next(iter_loader2)
            print('training head2')
    except StopIteration:
        # train remaining iterations from `iter_loader1` here?
        # Use `DataLoader`._num_yielded?
</code></pre>

rsomani95: <p>I think I figured it out. Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a>!</p>
<pre><code class="lang-python">dataset_head1 = TensorDataset(torch.randn(10, 1), torch.randn(10, 1))
dataset_head2 = TensorDataset(torch.randn(2, 1), torch.randn(2, 1))

heads = {
    'head1': dataset_head1,
    'head2': dataset_head2
}

for epoch in range(2):
    print('epoch ', epoch)
    print('========')
    # in epoch loop
    
    dls = {
        'head1': DataLoader(heads['head1']),
        'head2': DataLoader(heads['head2'])
    } 
    iters = {
        'head1': iter(dls['head1']),
        'head2': iter(dls['head2'])
    }
    
    try:
        while True:
            # Keep looping over dataloaders sequentially
            # until there are none for that head
            for name,iterator in iters.items():
                if not len(iterator) - iterator._num_yielded == 0:
                    data,target = next(iterator)
                    print(f'training {name}')
                else:
                    pass
            
            # raise `StopIteration` when all batches for all heads
            # have been yielded
            pending_per_head = [len(it) - it._num_yielded for _,it in iters.items()]
            if torch.tensor(pending_per_head).sum() == 0:
                raise StopIteration
    except StopIteration:
        pass
    print()
</code></pre>