Autoencoder with multiple decoders

Padmaksha_Roy: <p>Dear Team, I am trying to create an auto-encoder model which has one encoder and multiple decoders(depending on the number of classes). I am stuck while optimizing the loss. The loss is calculated by adding all the losses of the decoder output and the original data matrix for each class. I am not getting how to get it combinedly for all the classes together. Some insights will be helpful. Thanks</p>
<blockquote>
<pre><code>import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from dataset import Dataset
from autoencoder import Autoencoder
import numpy as np
</code></pre>
<p>device = torch.device(“cuda:0” if torch.cuda.is_available() else “cpu”)</p>
<p>num_epochs = 500<br>
batch_size = 50<br>
learning_rate = 0.01</p>
<p>model = Autoencoder().to(device)<br>
criterion = nn.MSELoss()<br>
class MultitaskAutoencoder(nn.Module):<br>
def <strong>init</strong>(self):</p>
<p>super(MultitaskAutoencoder,self).<strong>init</strong>()<br>
self.ec = nn.Linear(29,50)<br>
self.dc1= nn.Linear(50,29)<br>
self.dc2= nn.Linear(50,29)<br>
self.dc3 = nn.Linear(50,29)<br>
self.dc4= nn.Linear(50,29)<br>
self.dc5= nn.Linear(50,29)</p>
</blockquote>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/a/7/a77d25a47bf8c735f614c21d61dc5972367d859b.png" alt="Multi_AE" data-base62-sha1="nTFRyGXzkDcF3a1HLkec9TqnGUr" width="605" height="315"></p>
<blockquote>
<pre><code>def forward(self, x, dom=0):
    x = self.ec(x)
    out1 = self.dc1(x)
   # print ("out1:", out1.shape)

    out2 = self.dc2(x)
    #print ("out2:", out2.shape)

    out3 =self.dc3(x)
    #print ("out3:", out3.shape)
    #out4 = self.dc4(x)
    # out5 = self.dc5(x)
    final_out = torch.vstack((out1, out2, out3))

    return final_out
</code></pre>
<p>multitaskAE = MultitaskAutoencoder().to(device)</p>
<h1>print(multitaskAE)</h1>
<p>optimizer = torch.optim.Adam(multitaskAE.parameters(), lr=learning_rate, weight_decay=1e-5)</p>
<h1>test_data = torch.randn(3, 1, 256).to(device)</h1>
<h1>out = multitaskAE(test_data)</h1>
<h1>print(out.shape)</h1>
<p>def train_AE():<br>
dataset = Dataset()<br>
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)</p>
<pre><code>for epoch in range(100):
    train_loss = 0.0

    for input, out in dataloader:
        input = input.float().to(device)

        #print('input.shape = ', input.shape)  # 128 * 256
        dom_out = out.float().to(device) ## original data # 128 * 256
        #print('dom_out.shape = ', dom_out.shape) # 128 * 256

        output = multitaskAE(input)
        #print('predicted output.shape = ', output.shape) # 128 * 256 -- reconstructed values

        out1 = output[0:50,:]
        out2 = output[50:100,:]
        out3 = output[100:150,:]


        loss1 = criterion(out1, dom_out) 
        loss2 = criterion(out2,dom_out) # wrong 
        loss3 = criterion(out3,dom_out)

        final_loss = loss1 + loss2 + loss3

        final_loss.backward()
        optimizer.zero_grad()

        train_loss += final_loss.item()

    print('===&gt; Loss: ', train_loss / len(dataloader))
</code></pre>
<p>train_AE()</p>
</blockquote>

ptrblck: <p>Double post from <a href="https://discuss.pytorch.org/t/an-autoencoder-with-multiple-inputs/139138/8">here</a>.</p>