| import os |
| import torch |
| import torch.nn as nn |
| from transformers import PreTrainedModel, PretrainedConfig |
|
|
| class MVAEConfig(PretrainedConfig): |
| model_type = "mvae" |
|
|
| def __init__( |
| self, |
| prosoro_type="cylinder", |
| x_dim_dict=None, |
| h1_dim_dict=None, |
| h2_dim_dict=None, |
| z_dim=32, |
| layer_norm=False, |
| use_activation="relu", |
| **kwargs, |
| ): |
| super().__init__(**kwargs) |
| self.prosoro_type = prosoro_type |
| self.x_dim_dict = x_dim_dict |
| self.h1_dim_dict = h1_dim_dict |
| self.h2_dim_dict = h2_dim_dict |
| self.z_dim = z_dim |
| self.layer_norm = layer_norm |
| self.use_activation = use_activation |
|
|
|
|
| class MVAE(PreTrainedModel): |
| config_class = MVAEConfig |
|
|
| def __init__(self, config: MVAEConfig): |
| super().__init__(config) |
|
|
| self.prosoro_type = getattr(config, "prosoro_type", "cylinder") |
| self.x_dim_list = config.x_dim_dict |
| self.h1_dim_list = config.h1_dim_dict |
| self.h2_dim_list = config.h2_dim_dict |
| self.z_dim = config.z_dim |
|
|
| self.model = nn.ModuleDict() |
| self.model = nn.ModuleDict() |
| for i in range(len(self.x_dim_list)): |
| self.model[f"encoder_{i}"] = nn.Sequential( |
| nn.Linear(self.x_dim_list[i], self.h1_dim_list[i]), |
| nn.ReLU(), |
| nn.Linear(self.h1_dim_list[i], self.h2_dim_list[i]), |
| ) |
| self.model[f"encoder_{i}"] = nn.Sequential( |
| nn.Linear(self.x_dim_list[i], self.h1_dim_list[i]), |
| nn.ReLU(), |
| nn.Linear(self.h1_dim_list[i], self.h2_dim_list[i]), |
| ) |
| self.model[f"decoder_{i}"] = nn.Sequential( |
| nn.Linear(self.z_dim, self.h2_dim_list[i]), |
| nn.ReLU(), |
| nn.Linear(self.h2_dim_list[i], self.h1_dim_list[i]), |
| nn.ReLU(), |
| nn.Linear(self.h1_dim_list[i], self.x_dim_list[i]), |
| ) |
| self.model[f"decoder_{i}"] = nn.Sequential( |
| nn.Linear(self.z_dim, self.h2_dim_list[i]), |
| nn.ReLU(), |
| nn.Linear(self.h2_dim_list[i], self.h1_dim_list[i]), |
| nn.ReLU(), |
| nn.Linear(self.h1_dim_list[i], self.x_dim_list[i]), |
| ) |
| self.model[f"fc_mu_{i}"] = nn.Linear(self.h2_dim_list[i], self.z_dim) |
| self.model[f"fc_var_{i}"] = nn.Linear(self.h2_dim_list[i], self.z_dim) |
| self.model[f"fc_mu_{i}"] = nn.Linear(self.h2_dim_list[i], self.z_dim) |
| self.model[f"fc_var_{i}"] = nn.Linear(self.h2_dim_list[i], self.z_dim) |
|
|
| def sample(self, mu, var): |
| std = torch.exp(var / 2) |
| p = torch.distributions.Normal(torch.zeros_like(mu), torch.ones_like(std)) |
| q = torch.distributions.Normal(mu, std) |
| z = q.rsample() |
| return p, q, z |
|
|
| def x_to_z_encoder(self, x, input_index): |
| h = self.model[f"encoder_{input_index}"](x) |
| mu = self.model[f"fc_mu_{input_index}"](h) |
| var = self.model[f"fc_var_{input_index}"](h) |
| h = self.model[f"encoder_{input_index}"](x) |
| mu = self.model[f"fc_mu_{input_index}"](h) |
| var = self.model[f"fc_var_{input_index}"](h) |
| _, _, z = self.sample(mu, var) |
| return z |
|
|
| def z_to_x_decoder(self, z, output_index): |
| x_hat = self.model[f"decoder_{output_index}"](z) |
| x_hat = self.model[f"decoder_{output_index}"](z) |
| return x_hat |
|
|
| def forward(self, x): |
| x_hat_list = [] |
| for i in range(len(self.x_dim_list)): |
| z = self.x_to_z_encoder(x, 0) |
| x_hat = self.z_to_x_decoder(z, i) |
| x_hat_list.append(x_hat) |
| return x_hat_list[1:] |
|
|
| @classmethod |
| def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): |
| config = kwargs.get("config", None) |
|
|
| if config is None: |
| from transformers import AutoConfig |
| config = AutoConfig.from_pretrained(pretrained_model_name_or_path) |
| |
| prosoro_type = getattr(config, "prosoro_type", None) |
|
|
| pretrained_model_name_or_path = pretrained_model_name_or_path + f"/{prosoro_type}" |
|
|
| return super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs) |
|
|
