Spaces:
Running
Running
| import torch as th | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import numpy as np | |
| from geometry import relaxed_distortion_measure | |
| from pathlib import Path | |
| import json, os | |
| class AutoEncoderClass(nn.Module): | |
| def __init__( | |
| self, | |
| x_dim, | |
| z_dim=2, | |
| h_dims=[64,32], | |
| loss_type='l2', # 'l2', 'l1', 'huber' | |
| actv=nn.ReLU(), # activation function | |
| iso = True, | |
| iso_reg = 1e-7, | |
| ): | |
| super(AutoEncoderClass, self).__init__() | |
| self.name = 'autoencoder' | |
| self.z_dim = z_dim | |
| self.h_dims = h_dims | |
| self.loss_type = loss_type.lower() | |
| self.actv = actv | |
| self.iso = iso | |
| self.iso_reg = iso_reg | |
| if isinstance(x_dim, int): | |
| self.input_shape = (x_dim,) | |
| self.flattened_dim = x_dim | |
| elif isinstance(x_dim, tuple): | |
| self.input_shape = x_dim | |
| self.flattened_dim = int(np.prod(x_dim)) | |
| else: | |
| raise ValueError("x_dim must be int or tuple") | |
| enc_layers = [] | |
| in_dim = self.flattened_dim | |
| for h_dim in self.h_dims: | |
| enc_layers.append(nn.Linear(in_dim, h_dim)) | |
| enc_layers.append(self.actv) | |
| in_dim = h_dim | |
| enc_layers.append(nn.Linear(in_dim, self.z_dim)) | |
| self.encoder = nn.Sequential(*enc_layers) | |
| dec_layers = [] | |
| in_dim = self.z_dim | |
| for h_dim in reversed(self.h_dims): | |
| dec_layers.append(nn.Linear(in_dim, h_dim)) | |
| dec_layers.append(self.actv) | |
| in_dim = h_dim | |
| dec_layers.append(nn.Linear(in_dim, self.flattened_dim)) | |
| self.decoder = nn.Sequential(*dec_layers) | |
| def encode(self, x): return self.encoder(x.view(x.size(0), -1)) | |
| def decode(self, z): return self.decoder(z).view(z.size(0), *self.input_shape) | |
| def forward(self, x): z = self.encode(x); return self.decode(z), z | |
| def compute_loss(self, x, x_hat): | |
| if self.loss_type == 'l2': | |
| loss = nn.functional.mse_loss(x_hat, x) | |
| elif self.loss_type == 'l1': | |
| loss = nn.functional.l1_loss(x_hat, x) | |
| elif self.loss_type == 'huber': | |
| loss = nn.functional.huber_loss(x_hat, x) | |
| elif self.loss_type == 'mseloss': | |
| loss = nn.functional.mse_loss(x_hat,x) | |
| else: | |
| raise ValueError(f"Unknown loss_type: {self.loss_type}") | |
| if self.iso: | |
| iso_loss = relaxed_distortion_measure(self.decode, self.encode(x), eta=0.2, metric = 'identity') | |
| total_loss = loss + self.iso_reg * iso_loss | |
| else: | |
| iso_loss = th.tensor(0.0, device=x.device, dtype=loss.dtype) | |
| total_loss =loss | |
| return total_loss, loss, iso_loss | |
| def load_saved_model(device='cpu'): | |
| """ | |
| load AutoEncoderClass instance using saved config.json & model.pt | |
| Args: | |
| model_name | |
| device (str): 'cpu' / 'cuda:0' | |
| Returns: | |
| ae (AutoEncoderClass): model instance | |
| config (dict): configuration | |
| """ | |
| current_dir = Path(__file__).parent | |
| config_path = current_dir / "config.json" | |
| model_path = current_dir / "model.pt" | |
| if not config_path.exists(): | |
| raise FileNotFoundError(f"Cannot find path. Current path: {os.getcwd()}") | |
| with open(config_path, 'r') as f: | |
| config = json.load(f) | |
| actv_map = { | |
| 'Tanh': th.nn.Tanh(), | |
| 'ReLU': th.nn.ReLU(), | |
| 'SiLU': th.nn.SiLU(), | |
| 'GELU': th.nn.GELU(), | |
| } | |
| actv_fn = actv_map[config['actv']] | |
| ae = AutoEncoderClass( | |
| x_dim=config['x_dim'][0], | |
| z_dim=config['z_dim'], | |
| h_dims=config['h_dims'], | |
| loss_type=config['loss_type'], | |
| actv=actv_fn, | |
| iso=config['iso'], | |
| iso_reg=config['iso_reg'], | |
| ).to(device) | |
| ae.load_state_dict(th.load(model_path, map_location=device)) | |
| ae.eval() | |
| return ae, config |