0
$\begingroup$

I'm working on a Multi-Task VAE with one Encoder and two Decoders. The input consists of a vector with parameters which describe a design of a fluid system. The goal is to reconstruct the parameters with one decoder and predict the physical values of the design with the other one.

The model should be able to generate new designs with the correct physical values.

class MTVAE_CNN(nn.Module):
    def __init__(self, task_num, input_shape, input_size=28,z_dim=4, b=1,batch_size=10,c = 1, a = 1):
        super(MTVAE_CNN,self).__init__()
    
    
    self.task_num = task_num
    self.log_vars = nn.Parameter(torch.zeros((task_num)))
    
    self.z_dim = z_dim
    self.a = a
    self.b = b
    self.c = c
    self.input_size = input_size
    self.input_shape = input_shape
    self.batch_size = batch_size
    self.loss_list = []
    
    self.encoder = nn.Sequential(
        nn.ConvTranspose1d(1, 16, kernel_size=5, stride=3, padding=2),nn.ReLU(),nn.Dropout(p=0.2),
        nn.ConvTranspose1d(16, 32, kernel_size=5, stride=2, padding=2),nn.ReLU(),nn.Dropout(p=0.2),
        nn.Conv1d(32, 64, kernel_size=5, stride=2),nn.ReLU(),nn.Dropout(p=0.2),
        nn.MaxPool1d(7),nn.ReLU(),nn.Dropout(p=0.2),
        nn.Flatten()
        )


    self.decoder = nn.Sequential(
        nn.Unflatten(1,(1,self.z_dim)),
        nn.ConvTranspose1d(in_channels=1, out_channels=16, kernel_size=5, stride=3),nn.ReLU(),
        nn.Conv1d(16, 32, kernel_size=5, stride=3),nn.ReLU(),
        nn.MaxPool1d(7),nn.ReLU(),nn.Dropout(p=0.2),
        nn.Flatten()
    )

    self.predictor = nn.Sequential(
        nn.Unflatten(1,(1,self.z_dim)),
        nn.ConvTranspose1d(1, 16, kernel_size=5, stride=3),nn.ReLU(),
        nn.Conv1d(16, 32, kernel_size=5, stride=3),nn.ReLU(),
        nn.MaxPool1d(3),nn.ReLU(),nn.Dropout(p=0.2),
        nn.Flatten()
    )


    self.conv_out_size = self.calculate_conv_output_size(self.input_shape)

    
    self.Bn1 = nn.Linear(self.conv_out_size, self.z_dim)
    self.Bn2 = nn.Linear(self.conv_out_size, self.z_dim)

    self.dec_lin = nn.Linear(self.calculate_decoder_output_size(self.input_shape),self.input_size)
    self.pre_lin = nn.Linear(self.calculate_predictor_output_size(self.input_shape),2)

    self.sig1 = nn.ReLU()
    self.sig2 = nn.Sigmoid()

def calculate_conv_output_size(self,input):

    x = torch.randn(*input) # Create a random input tensor
    x = x.unsqueeze(1)
    x = self.encoder(x)
    
    return x.view(x.size(0), -1).size(1)

def calculate_decoder_output_size(self,input):

    x = torch.randn(*input) # Create a random input tensor
    x = x.unsqueeze(1)
    x = self.encoder(x)
    x = self.Bn1(x)
    x = self.decoder(x)
    
    
    return x.view(x.size(0), -1).size(1)

def calculate_predictor_output_size(self,input):

    x = torch.randn(*input) # Create a random input tensor
    x = x.unsqueeze(1)
    x = self.encoder(x)
    x = self.Bn1(x)
    x = self.predictor(x)
    
    return x.view(x.size(0), -1).size(1)

def reparameterise(self, mu, logvar):
    if self.training:
        std = logvar.mul(0.5).exp_()
        eps = std.data.new(std.size()).normal_()
        return eps.mul(std).add_(mu)
    else:
        return mu
    
def bottleneck(self,h):
    
    mu,logvar = self.Bn1(h), self.Bn2(h)
       
    z = self.reparameterise(mu,logvar)
    
    return z, mu, logvar 
    
def predict (self,z):

    #y1 = self.predictor(torch.cat([z.unsqueeze(-1), y.unsqueeze(-1)], dim=1))
    
    y1 = self.predictor(z)

    y1 = self.pre_lin(y1)

    y1 = self.sig1(y1)
         
    return y1

def encode(self,x):
    
    x = self.encoder(x.unsqueeze(1))
    #mu = mu_logvar[:, 0, :]
    #logvar = mu_logvar[:, 1, :]

    #x = x.view(x.size()[0], -1)

    mu,logvar = self.Bn1(x), self.Bn2(x)

    z = self.reparameterise(mu,logvar)

    return z,mu,logvar

def decode(self, z):
    
    y = self.decoder(z)

    y = self.dec_lin(y)

    y = self.sig2(y)

    
    
    return y

def forward(self, x):

    z, mu, logvar = self.encode(x)
    y_hat = self.predict(z)
    x_hat = self.decode(z)
    
    return x_hat, y_hat, mu, logvar

def loss_function(self, x,y,b):
    
    x_hat, y_hat, mu, logvar = self.forward(x)
    
    MSE_Reg = nn.MSELoss(reduction = "sum")(y_hat, y)
    
    MSE = nn.MSELoss(reduction = "sum")(
        x_hat, x)
    
    KLD = 0.5 * torch.sum(logvar.exp() - logvar - 1 + mu.pow(2))
    
    kld_loss = b*KLD
    
    loss1 = self.a*(MSE)
    
    loss2 = self.c*MSE_Reg
    
    return loss2,loss1,kld_loss,b

I tried different optimization techniques for multitask learning, grid search for the loss weights, uncertainty weighted loss and increasing the hidden layers of the encoder. The 1DConvolution Layers have been delivering the best results, but theres still room for improvement.

The model should concentrate on representing the designs in the latent space with the information about the values, without focusing too much on the representation of the values.

Is there a way to influence the behavior and the performance of the model with the last decoder activation functions?

$\endgroup$

0

You must log in to answer this question.