I have the distribution of each weight (and bias) of a neural network. What is a fast way to generate sample neural networks? To be specific, the distribution is assumed to be Gaussian and the mean and variance are stored in two dictionaries.
I have a small ResNet with about 400,000 parameters (weights and biases). The fastest way I got so far takes about 1 minute to generate a sample neural network on a GPU. Any ideas how to speed it up?
My code is as follows. It iterates through individual weights and biases to generate a sample of a scalar Gaussian random variable at a time.
I also tried generating a vector (the same size as that of param
) at a time by using torch.distributions.MultivariateNormal and diag() matrices, but I got an out-of-memory error.
def generate_net(dic1, dic2, sample_net):
# dic1: dictionary storing mean
# dic2: dictionary storing variance
# sample_net: model of a network
for name, param in sample_net.named_parameters():
param1 = dic1[name]
param2 = dic2[name]
variance_final = param2 - torch.mul(param1, param1)
size_tmp = param2.size()
param1 = param1.view(-1, 1) # make it a ?x1 vector
variance_final = variance_final.view(-1, 1)
param = param.view(-1,1)
indx = 0
for wt1, var in zip(param1, variance_final):
rv = torch.distributions.Normal(wt1, torch.sqrt(var) * damp_factor)
param[indx, 0] = rv.sample()
param = param.view(size_tmp) # get back to original size
param1 = param1.view(size_tmp)
param2 = param2.view(size_tmp)
Thanks.