def train_vae_mlp(latent_code_num, hidden, params, device):
    print((('Totally ' + str(params['VAE_epoch_num'])) + ' epochs to train'))
    os.environ['CUDA_VISIBLE_DEVICES'] = str(device)
    thermal_conductivity_train_loader = torch.load('Data/thermal_conductivity_vae_mlp_train_loader.pkl')
    heat_capacity_train_loader = torch.load('Data/heat_capacity_vae_mlp_train_loader.pkl')
    heat_capacity_vae_mlp = VAE_MLP_CAT(latent_code_num, hidden).cuda()
    thermal_conductivity_vae_mlp = VAE_MLP_CAT(latent_code_num, hidden).cuda()
    thermal_conductivity_optimizer = optim.Adam(thermal_conductivity_vae_mlp.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0)
    heat_capacity_optimizer = optim.Adam(heat_capacity_vae_mlp.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0)
    thermal_conductivity_total_loss_list = np.ones((params['VAE_epoch_num'] + 10))
    heat_capacity_total_loss_list = np.ones((params['VAE_epoch_num'] + 10))
    thermal_conductivity_total_loss_list *= 5000000
    heat_capacity_total_loss_list *= 5000000
    thermal_conductivity_model_file_name = (((('Model_pkl/VAE_MLP_CAT_thermal_conductivity_latent_' + str(latent_code_num)) + '_structure_') + str(hidden)) + '.pkl')
    heat_capacity_model_file_name = (((('Model_pkl/VAE_MLP_CAT_heat_capacity_latent_' + str(latent_code_num)) + '_structure_') + str(hidden)) + '.pkl')
    for epoch in range(params['VAE_epoch_num']):
        total_loss = 0
        thermal_conductivity_vae_mlp.train()
        for (i, data) in enumerate(tqdm(thermal_conductivity_train_loader, 0)):
            one_hot = torch.cat((data[0], data[1]), dim=1)
            one_hot = one_hot.reshape(one_hot.shape[0], 1, one_hot.shape[1], one_hot.shape[2])
            one_hot = Variable(one_hot).cuda().type(torch.cuda.FloatTensor)
            thermal_conductivity_optimizer.zero_grad()
            t = data[2].cuda().reshape(data[2].shape[0], 1).type(torch.cuda.FloatTensor)
            label = data[3].cuda().reshape(data[3].shape[0], 1).type(torch.cuda.FloatTensor)
            (recon_x, mu, logvar, pre) = thermal_conductivity_vae_mlp.forward(one_hot, t)
            recon_x = recon_x[:, :, :one_hot.shape[2], :one_hot.shape[3]]
            loss = loss_func(recon_x, one_hot, mu, logvar, pre, label)
            loss.backward()
            total_loss += (loss.data.item() / 1000)
            thermal_conductivity_optimizer.step()
        print('====> Epoch: {} Average loss: {:.4f}'.format(epoch, (total_loss / len(thermal_conductivity_train_loader.dataset))))
        thermal_conductivity_total_loss_list[epoch] = (total_loss / len(thermal_conductivity_train_loader.dataset))
        if (np.argmin(thermal_conductivity_total_loss_list) == epoch):
            torch.save(heat_capacity_vae_mlp, thermal_conductivity_model_file_name)
            print(('best result, saving the model to ' + thermal_conductivity_model_file_name))
        elif (np.argmin(thermal_conductivity_total_loss_list) == (epoch - 25)):
            print('Finish: Training process over due to useless training')
            break
    for epoch in range(params['VAE_epoch_num']):
        total_loss = 0
        heat_capacity_vae_mlp.train()
        for (i, data) in enumerate(tqdm(heat_capacity_train_loader, 0)):
            one_hot = torch.cat((data[0], data[1]), dim=1)
            one_hot = one_hot.reshape(one_hot.shape[0], 1, one_hot.shape[1], one_hot.shape[2])
            one_hot = Variable(one_hot).cuda().type(torch.cuda.FloatTensor)
            heat_capacity_optimizer.zero_grad()
            t = data[2].cuda().reshape(data[2].shape[0], 1).type(torch.cuda.FloatTensor)
            (recon_x, mu, logvar, pre) = heat_capacity_vae_mlp.forward(one_hot, t)
            recon_x = recon_x[:, :, :one_hot.shape[2], :one_hot.shape[3]]
            loss = loss_func(recon_x, one_hot, mu, logvar, pre, t)
            loss.backward()
            total_loss += (loss.data.item() / 1000)
            heat_capacity_optimizer.step()
        print('====> Epoch: {} Average loss: {:.4f}'.format(epoch, (total_loss / len(heat_capacity_train_loader.dataset))))
        heat_capacity_total_loss_list[epoch] = (total_loss / len(heat_capacity_train_loader.dataset))
        if (np.argmin(heat_capacity_total_loss_list) == epoch):
            torch.save(heat_capacity_vae_mlp, heat_capacity_model_file_name)
            print(('best result, saving the model to ' + heat_capacity_model_file_name))
        elif (np.argmin(thermal_conductivity_total_loss_list) == (epoch - 25)):
            print('Finish: Training process over due to useless training')
            break
def train_vae_mlp(latent_code_num, hidden, params, device):
    print((('Totally ' + str(params['VAE_epoch_num'])) + ' epochs to train'))
    os.environ['CUDA_VISIBLE_DEVICES'] = str(device)
    thermal_conductivity_train_loader = torch.load('Data/thermal_conductivity_vae_mlp_train_loader.pkl')
    heat_capacity_train_loader = torch.load('Data/heat_capacity_vae_mlp_train_loader.pkl')
    heat_capacity_vae_mlp = VAE_MLP_CAT(latent_code_num, hidden).cuda()
    thermal_conductivity_vae_mlp = VAE_MLP_CAT(latent_code_num, hidden).cuda()
    thermal_conductivity_optimizer = optim.Adam(thermal_conductivity_vae_mlp.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0)
    heat_capacity_optimizer = optim.Adam(heat_capacity_vae_mlp.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0)
    thermal_conductivity_total_loss_list = np.ones((params['VAE_epoch_num'] + 10))
    heat_capacity_total_loss_list = np.ones((params['VAE_epoch_num'] + 10))
    thermal_conductivity_total_loss_list *= 5000000
    heat_capacity_total_loss_list *= 5000000
    thermal_conductivity_model_file_name = (((('Model_pkl/VAE_MLP_CAT_thermal_conductivity_latent_' + str(latent_code_num)) + '_structure_') + str(hidden)) + '.pkl')
    heat_capacity_model_file_name = (((('Model_pkl/VAE_MLP_CAT_heat_capacity_latent_' + str(latent_code_num)) + '_structure_') + str(hidden)) + '.pkl')
    for epoch in range(params['VAE_epoch_num']):
        total_loss = 0
        thermal_conductivity_vae_mlp.train()
        for (i, data) in enumerate(tqdm(thermal_conductivity_train_loader, 0)):
            one_hot = torch.cat((data[0], data[1]), dim=1)
            one_hot = one_hot.reshape(one_hot.shape[0], 1, one_hot.shape[1], one_hot.shape[2])
            one_hot = Variable(one_hot).cuda().type(torch.cuda.FloatTensor)
            thermal_conductivity_optimizer.zero_grad()
            t = data[2].cuda().reshape(data[2].shape[0], 1).type(torch.cuda.FloatTensor)
            label = data[3].cuda().reshape(data[3].shape[0], 1).type(torch.cuda.FloatTensor)
            (recon_x, mu, logvar, pre) = thermal_conductivity_vae_mlp.forward(one_hot, t)
            recon_x = recon_x[:, :, :one_hot.shape[2], :one_hot.shape[3]]
            loss = loss_func(recon_x, one_hot, mu, logvar, pre, label)
            loss.backward()
            total_loss += (loss.data.item() / 1000)
            thermal_conductivity_optimizer.step()
        print('====> Epoch: {} Average loss: {:.4f}'.format(epoch, (total_loss / len(thermal_conductivity_train_loader.dataset))))
        thermal_conductivity_total_loss_list[epoch] = (total_loss / len(thermal_conductivity_train_loader.dataset))
        if (np.argmin(thermal_conductivity_total_loss_list) == epoch):
            torch.save(heat_capacity_vae_mlp, thermal_conductivity_model_file_name)
            print(('best result, saving the model to ' + thermal_conductivity_model_file_name))
        elif (np.argmin(thermal_conductivity_total_loss_list) == (epoch - 25)):
            print('Finish: Training process over due to useless training')
            break
    for epoch in range(params['VAE_epoch_num']):
        total_loss = 0
        heat_capacity_vae_mlp.train()
        for (i, data) in enumerate(tqdm(heat_capacity_train_loader, 0)):
            one_hot = torch.cat((data[0], data[1]), dim=1)
            one_hot = one_hot.reshape(one_hot.shape[0], 1, one_hot.shape[1], one_hot.shape[2])
            one_hot = Variable(one_hot).cuda().type(torch.cuda.FloatTensor)
            heat_capacity_optimizer.zero_grad()
            t = data[2].cuda().reshape(data[2].shape[0], 1).type(torch.cuda.FloatTensor)
            (recon_x, mu, logvar, pre) = heat_capacity_vae_mlp.forward(one_hot, t)
            recon_x = recon_x[:, :, :one_hot.shape[2], :one_hot.shape[3]]
            loss = loss_func(recon_x, one_hot, mu, logvar, pre, t)
            loss.backward()
            total_loss += (loss.data.item() / 1000)
            heat_capacity_optimizer.step()
        print('====> Epoch: {} Average loss: {:.4f}'.format(epoch, (total_loss / len(heat_capacity_train_loader.dataset))))
        heat_capacity_total_loss_list[epoch] = (total_loss / len(heat_capacity_train_loader.dataset))
        if (np.argmin(heat_capacity_total_loss_list) == epoch):
            torch.save(heat_capacity_vae_mlp, heat_capacity_model_file_name)
            print(('best result, saving the model to ' + heat_capacity_model_file_name))
        elif (np.argmin(thermal_conductivity_total_loss_list) == (epoch - 25)):
            print('Finish: Training process over due to useless training')
            break