Hi, I’m building a generator g,  that receives a latent-code (vector of shape 100) and outputs an image. Specifically, I have 1000 MNIST images, and I want the network to learn a latent code z_i for each image x_i, such that g(z_i)=x_i (this approach is known as Generative Latent Optimization). So I’ve used nn.Embedding(1000,embedding_dim=100) and standard generator architecture, that receives the code from embedding and outputs an image. As for loss, I combine reconstruction loss with regularization on the embedding-vector weights.
My probelm is: I’d like to add noise to the latent-code vector before it is inserted to the generator (in order to make the latent-code compact). However I’m a beginner, and I don’t know whether I should call detach() when adding the noise or not. I’m not sure of my approach entirely. I don’t want to learn the scale of the noise or anything.
Here’s my attempt:
class net(nn.Module):
  def __init__():
    self.embed = nn.Embedding(1000,embedding_dim=100)
    self.generator = nn.sequential( nn.Linear(100, 84), .... )
  def forward(batch_indices):
    batch_codes = self.embed(batch_indices)
    noise = torch.randn_like(batch_codes) * sigma
    noisy_batch_codes = batch_codes + noise # SHOULD THIS BE batch_codes.detach() + noise ??
    return self.generator(noisy_batch_codes)

g = net()
optim = SGD(g.parameters(), lr=0.01)
for epoch in range(num_epochs):
  for orig_images, orig_images_idx in trainloader:
    optim.zero_grad()
    output = g(orig_images_idx)
    reconstruction_loss = nn.MSELoss()(output, orig_images)
    embed_vector_weights = g.embed.weight[orig_images_idx]
    reg_loss = torch.norm(embed_vector_weights) * reg_coeff
    loss = reconstruction_loss + reg_loss
    loss.backward()
    optim.step()