Multi-task autoencoder loss stuck and not converging

username: Padmaksha_Roy
cooked: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a> , I am trying to implement a multi-task loss function where I have one encoder and multiple decoders. The input is from multiple domains and it tries to jointly achieve good reconstruction of all source data given a particular segment of train data.</p>
<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>
<p>In order to achieve it, all the samples from different domains are concatenated together like this. Here 1-M are the different domains and l is the replicated data from the lth domain.<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/c/0/c001b18963e611db5d968c7cfeec8b09c630ee89.png" alt="image" data-base62-sha1="roziLlkwqlw5NfovOccDEOmD44p" width="354" height="86"></p>
<p>The code is below:</p>
<blockquote>
<p>import torch<br>
import torch.nn as nn<br>
from torch.utils.data import DataLoader<br>
from dataset import Dataset<br>
from autoencoder import Autoencoder<br>
import numpy as np</p>
<p>device = torch.device(“cuda:0” if torch.cuda.is_available() else “cpu”)</p>
<p>num_epochs = 500<br>
batch_size = 27<br>
learning_rate = 0.005<br>
feats = 29<br>
domains = 3<br>
latent_dims = 10</p>
<p>model = Autoencoder().to(device)<br>
criterion = nn.MSELoss()</p>
<p>class MultitaskAutoencoder(nn.Module):</p>
<pre><code>def __init__(self):
    super().__init__()

    self.encoder = nn.Sequential(
        nn.Linear(feats, latent_dims),
        nn.ReLU(),
        # nn.Linear(4, 3)
    )
    self.decoders = nn.ModuleList()
    for _ in range(batch_size):
        self.decoders.append(nn.Linear(10, feats))

    self.encoder_hidden_layer = nn.Linear(
        in_features=feats, out_features=20
    )
    self.encoder_output_layer = nn.Linear(
        in_features=20, out_features=10
    )


def forward(self, inputs): #inputs--&gt; batch_size

    activation = self.encoder_hidden_layer(inputs)
    activation = torch.relu(activation)
    code = self.encoder_output_layer(activation)
    code = torch.relu(code)
    #z.append(code)

    outs = []
    for idx, dec in enumerate(self.decoders):
        outs.append(dec(code[idx]))
    outs = torch.stack(outs)
    return outs
</code></pre>
<p>multitaskAE = MultitaskAutoencoder().to(device)<br>
optimizer = torch.optim.Adam(multitaskAE.parameters(), lr=learning_rate, weight_decay=1e-5)</p>
<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

        loss = criterion(output, dom_out)

        loss.backward()
        optimizer.zero_grad()

        train_loss += loss.item()

    print('===&gt; Loss: ', train_loss / len(dataloader))
</code></pre>
<p>train_AE()</p>
</blockquote>
<p>The dataset generator is as below:</p>
<blockquote>
<p>def construct_pair(X_list):  # 3 * 10000 * 29</p>
<pre><code>n_dom = len(X_list) # 3 * 10000 * 29
#print("Xin shape BEFORE virtual stack:", X_list.shape)

X_in = np.vstack(X_list)  ## 30000 * 29
print("Xin SHAPE TRAIN DATA", X_in.shape)

X_outs = [] # ground truth data
for i in range(0, n_dom): # for each domain
    X = X_list[i]  # take first domain 1 * 10000 * 29

    print("X shape:", X.shape) ## X shape: 1 * 10000 * 29

    Z_list = []
    for j in range(0, n_dom):
        Z_list.append(X) # make 3 (num of domains) copies of the same 1000 samples. 3 * 10000 * 29

    Z = np.vstack(Z_list) # Z shape: 30000 * 29
    print("Z shape:", Z.shape)

    X_outs.append(Z) #  3 * 30000 * 29 ## same samples repeated to nos of domains.

print("XOUT SHAPE ", len(X_outs))
return X_in, X_outs
</code></pre>
</blockquote>
<blockquote>
<p>Blockquote</p>
</blockquote>
<p>My loss value is stuck and loss is not converging. Can you please give some hints on what am I doing wrong here? Thanks a lot!</p>

username: Padmaksha_Roy
cooked: <p>Hi, <a class="mention" href="/u/ptrblck">@ptrblck</a> , can you plz provide me some hints here? Thanks!</p>

username: ptrblck
cooked: <p>I’ve responded to your cross-post already and mentioned that I wouldn’t know which approach works best for your use case (<a href="https://discuss.pytorch.org/t/an-autoencoder-with-multiple-inputs/139138/11">here</a>). I would thus recommend avoiding tagging specific users as it might discourage others to post a valid answer.<br>
Generally, you could try to overfit a small dataset (e.g. just 10 samples) and make sure your model is able to learn these samples (almost) perfectly. If that’s not the case your model architecture or overall training could have other errors or might just not be suitable for the task.</p>

username: Padmaksha_Roy
cooked: <p>hi , I am sorry for this tagging mistake, and thanks for your suggestion. I will keep this open now so that I can get some hints from others as well.</p>

username: Padmaksha_Roy
cooked: <p>Hi All, the idea I am trying to implement is from this paper “Domain Generalization for Object Recognition with Multi-task Autoencoders”-  <a href="https://arxiv.org/pdf/1508.07680.pdf" rel="noopener nofollow ugc">https://arxiv.org/pdf/1508.07680.pdf</a>.<br>
I have a confusion about the decoder architecture - do I need to split the encoder output into M parts and pass it through the decoder with equal inputs? Shall i keep all decoder architectures the same here? Kindly provide some input here .</p>