How to use nn.init.kaiming_uniform to initialize the weights of my VAE?

PyTorch2603: <p>Hi. I created a VAE which takes in a 720x1 vector of numbers as the input. I wanted to know how I could initialize the weights of my VAE using kaiming_uniform, so I could improve its performance. I tried searching online but was unable to put this in place.<br>
Here is the code:<br>
import os<br>
import numpy as np<br>
import matplotlib.pyplot as plt<br>
import numpy as np<br>
import torch<br>
import torch.nn as nn<br>
import torch.nn.functional as F<br>
import torch.optim as optim<br>
import torchvision<br>
from torchvision import datasets, transforms<br>
from torch.autograd import Variable<br>
from torchvision.utils import save_image<br>
import os<br>
from PIL import Image as PImage<br>
device=“cpu”<br>
trainloss=0<br>
testloss=0<br>
class VAE(nn.Module):<br>
def <strong>init</strong>(self, z_dim):<br>
super(VAE, self).<strong>init</strong>()</p>
<pre><code>    # Encoder
    self.encode=nn.Sequential(
    nn.Linear(720,540),
    nn.ReLU(),
    nn.Linear(540,400),
    nn.ReLU(),
    nn.Linear(400,300),
    nn.ReLU(),
    nn.Linear(300,225),
    nn.ReLU(),
    nn.Linear(225,150),
    nn.ReLU(),
    nn.Linear(150,100),
    nn.ReLU(),
    nn.Linear(100,60),
    nn.ReLU(),
    nn.Linear(60,45)
    )
    self.fc_mu = nn.Linear(45, z_dim)
    self.fc_logvar = nn.Linear(45, z_dim)
    # Decoder
    self.decode=nn.Sequential(
    nn.Linear(z_dim, 45),
    nn.ReLU(),
    nn.Linear(45, 60),
    nn.ReLU(),
    nn.Linear(60,100),
    nn.ReLU(),
    nn.Linear(100,150),
    nn.ReLU(),
    nn.Linear(150,225),
    nn.ReLU(),
    nn.Linear(225,300),
    nn.ReLU(),
    nn.Linear(300,400),
    nn.ReLU(),
    nn.Linear(400,540),
    nn.ReLU(),
    nn.Linear(540,720)
    )
    
def encoder(self, x):
    a=self.encode(x)
    return self.fc_mu(a),self.fc_logvar(a)
    

def sampling(self, mu, logvar):
    std = torch.exp(0.5 * logvar)
    eps = torch.randn_like(std)
    return eps.mul(std).add_(mu) # return z sample

def decoder(self, z):
    return self.decode(z)

def forward(self, x):
    mu, logvar = self.encoder(x)
    z = self.sampling(mu, logvar)
    return self.decoder(z), mu, logvar
</code></pre>
<p>vae = VAE(z_dim=20).to(device)<br>
optimizer = optim.Adam(vae.parameters(), lr=0.001)<br>
def loss_function(recon_x, x, mu, logvar):<br>
myloss=nn.MSELoss()<br>
MSE=myloss(recon_x,x)<br>
KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())<br>
return MSE + KLD<br>
def trainortest(myvector):<br>
optimizer.zero_grad()<br>
recon_batch, mu, logvar = vae(myvector)<br>
loss = loss_function(recon_batch, myvector, mu, logvar)<br>
if(epoch&lt;34000):<br>
loss.backward()<br>
myloss=loss.item()<br>
optimizer.step()<br>
else:<br>
myloss=loss.item()<br>
print(myloss)<br>
if(epoch==100):<br>
print(input_list)<br>
print(recon_batch)<br>
return myloss<br>
with open(‘x_y (1).txt’, ‘r’) as f:<br>
lines = f.readlines()<br>
epoch=0</p>
<p>trainloss=0<br>
testloss=0<br>
for line in lines:<br>
epoch+=1<br>
x, y = line.split(‘=’)[0], line.split(‘=’)[1]<br>
x, y = x.split(’ ‘), y.split(’ ')<br>
x = [i for i in x if i]<br>
y = [i for i in y if i]<br>
x[0] = x[0][1:]<br>
y[0] = y[0][1:]<br>
x[-1] = x[-1][:-1]<br>
y[-1] = y[-1][:-1]</p>
<pre><code>    x = [float(i) for i in x if i]
    y = [float(i) for i in y if i]
    S=np.zeros(360, dtype='complex_')
    for i in range(0,360):
        for k in range(360):
            a=x[k]
            b=y[k]
            tmp = ((-2j*np.pi*i*k)) /360
            S[i] += (complex(a,b)) * np.exp(tmp)
        S[i]=S[i]/360
    input_list=[]
    for i in range(0,360):
        input_list.append(np.real(S[i]))
        input_list.append(np.imag(S[i]))
    input_list=torch.FloatTensor(input_list)
    answer=trainortest(input_list)
    if(epoch&lt;34000):
        trainloss+=answer
    else:
        testloss+=answer
</code></pre>
<p>I would appreciate any help as to how I could initialize the weights of my VAE in the above code. Thank you!</p>

ptrblck: <p>You can define a custom <code>weights_init</code> method and <code>.apply</code> it as seen here:</p>
<pre><code class="lang-python">def weights_init(m):
    if isinstance(m, nn.Linear):
        torch.nn.init.kaiming_uniform_(m.weight)
        
model = nn.Linear(10, 10)
model.apply(weights_init)
</code></pre>

PyTorch2603: <p>Oh okay, thank you for the response!</p>