How to properly get the latent vector of a VAE during the traning?

username: DidierDeschamps
cooked: <p>Hi all,</p>
<p>I have a question, I found a solution to brutforce the problem but that’s not very pretty coding. I was wondering if you had a better a solution for my following problem.</p>
<p>I have made a VAE that’s in the following form:</p>
<pre><code class="lang-auto">def forward(self, input):
    latent_vector = encoder(input)
    output = decoder(latent_vector)
    return output 
</code></pre>
<p>I actually have loss in the following form:</p>
<pre><code class="lang-auto">predicted = self.model(inputs).float()
loss = self.criterion(predicted, expected)
loss.backward()
self.optimizer.step()
</code></pre>
<p>I would like to force the latent vector to have a particular form, so I would like to add a term in the loss like:</p>
<pre><code class="lang-auto">loss = self.criterion(predicted, expected) + some_function(latent_vector)
</code></pre>
<p>Is there a proper way to get my latent vector to perform that ?</p>
<p>Best,<br>
Didier</p>

username: Nikronic
cooked: <p>Hi,</p>
<p>If I have understood your question properly, You can just return <code>output</code> and <code>latent_vector</code> and do the other things as you have mentioned. But if you do not want to return the <code>output</code> not a tuple, I think it is good to add a attribute to your <code>Model</code> class and update it every time through <code>forward</code> function and you can obtain its value just by calling <code>model.self.latent_vector</code>.</p>
<pre><code class="lang-python">class Model:
  def __init__(self):
    self.latent_vector = None

  def forward(self, x):
    self.latent_vector = x
    output = x*2
    return output

model(2)  # prints 4
model.latent_vector  # prints 2
</code></pre>
<p>Good luck<br>
Nik</p>

username: DidierDeschamps
cooked: <p>Okay thank you so much for your answer in fact I just found a problem with this method according to my problem<br>
In fact I don’t want to look at the latent vector of the current input.</p>
<p>I have different classes let say for examples animals : cats, dogs and birds.<br>
If my current input is a cat: cat_i and I have n cats in my dataset: {cat_1, cat_2, …, cat_n}.<br>
I want the latent vector to learn the distribution among the cats, the dogs and the birds separately.</p>
<p>So I would like to code the Loss this way:</p>
<pre><code class="lang-auto">Loss = self.criterion(current_reconstructed_cat_i, cat_i) 
+ sum{for j=1 to n, j!=i} some_function(current_latent_vector_of_cat_j, current_latent_vector_of_cat_i)
</code></pre>
<p>current_latent_vector_of_cat_j is the current pass forward stopped to the latent_vector.</p>
<p>Best,<br>
Didier</p>

username: Nikronic
cooked: <p>So based on your explanation, you need all latent vectors of the dataset to be able to calculate the for loop. Actually, it may be tricky, because of memory issues because you need to pass all the dataset containing a class such as <code>cat</code>.</p>
<p>There is point here. the <code>encoder</code> generates <code>latent_vector</code>, so if you want the latent vector of whole dataset to calculate loss for only a batch, you have to prepare latent vectors one the go. I mean when you passed a batch through the network to calculate below code:</p>
<aside class="quote no-group" data-username="DidierDeschamps" data-post="3" data-topic="53016">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/didierdeschamps/48/12798_2.png" class="avatar"> DidierDeschamps:</div>
<blockquote>
<p>Loss = self.criterion(current_reconstructed_cat_i, cat_i)</p>
</blockquote>
</aside>
<p>you need all latent vectors, then you will optimize your model a <code>step</code>. So for the next batch, the latent vectors you already have are no longer valid and you have to calculate them again. This will be a slow process.But let say there is no such a problem. To calculate latent vector without updating weights, you can do following in your training procedure:</p>
<pre data-code-wrap="python"><code class="lang-python">latent_vectors = []
with torch.no_grad():
        for data in data_loader:
            latent_vector = model.encoder(data)
            latent_vectors.append(latent_vector)
 return latent_vectors
</code></pre>
<p>The only point is to use <code>torch.no_grad</code> to preserve the models state. But still I am concerned about the performance.</p>