Spaces:
Sleeping
Sleeping
| # Class to define the network architecture of the models | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from torch.utils.data import DataLoader | |
| from torch.optim import Adam | |
| class VanillaLSTM(nn.Module): | |
| def __init__( | |
| self, input_dim=1, hidden_dim=64, output_dim=1, num_layers=2, dropout=0.2 | |
| ): | |
| super(VanillaLSTM, self).__init__() | |
| self.hidden_dim = hidden_dim | |
| self.num_layers = num_layers | |
| self.lstm = nn.LSTM( | |
| input_size=input_dim, | |
| hidden_size=hidden_dim, | |
| num_layers=num_layers, | |
| batch_first=True, | |
| dropout=dropout, | |
| ) | |
| self.fc = nn.Linear(in_features=hidden_dim, out_features=output_dim) | |
| def forward(self, x): | |
| h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).requires_grad_() | |
| c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).requires_grad_() | |
| out, (hn, cn) = self.lstm(x, (h0.detach(), c0.detach())) | |
| out = self.fc(out[:, -1, :]) | |
| return out | |
| class VAE(nn.Module): | |
| def __init__(self, seq_len=48, n_features=1, hidden_dim=64, latent_dim=16, dropout=0.3): | |
| super(VAE, self).__init__() | |
| self.seq_len = seq_len | |
| self.hidden_dim = hidden_dim | |
| # Encoder | |
| self.enc_lstm = nn.LSTM( | |
| input_size=n_features, | |
| hidden_size=hidden_dim, | |
| batch_first=True | |
| ) | |
| self.enc_dropout = nn.Dropout(p=dropout) | |
| self.fc_mu = nn.Linear(hidden_dim, latent_dim) | |
| self.fc_var = nn.Linear(hidden_dim, latent_dim) | |
| # Decoder | |
| self.fc_upsample = nn.Linear(latent_dim, seq_len * hidden_dim) | |
| self.dec_dropout = nn.Dropout(p=dropout) | |
| self.dec_lstm = nn.LSTM( | |
| input_size=hidden_dim, | |
| hidden_size=hidden_dim, | |
| batch_first=True | |
| ) | |
| self.fc_out = nn.Linear(hidden_dim, n_features) | |
| def reparameterize(self, mu, log_var): | |
| std = torch.exp(0.5 * log_var) | |
| eps = torch.randn_like(std) | |
| return mu + eps * std | |
| def forward(self, x): | |
| # Encode | |
| _, (h_enc, c_enc) = self.enc_lstm(x) | |
| h_enc = h_enc.squeeze(0) # shape: (batch_size, hidden_dim) | |
| h_enc = self.enc_dropout(h_enc) | |
| mu, log_var = self.fc_mu(h_enc), self.fc_var(h_enc) | |
| # Reparameterize at latent space | |
| z = self.reparameterize(mu, log_var) | |
| # Decode | |
| z = self.fc_upsample(z) | |
| z = z.view(-1, self.seq_len, self.hidden_dim) | |
| decoded, _ = self.dec_lstm(z) | |
| dec_out = self.dec_dropout(decoded) | |
| out = self.fc_out(dec_out) | |
| return out, mu, log_var | |
| class Transformer(nn.Module): | |
| def __init__(self, input_dim=1, model_dim=64, num_layers=2, num_heads=4, dropout=0.2): | |
| super(Transformer, self).__init__() | |
| self.model_dim = model_dim | |
| self.num_layers = num_layers | |
| self.embedding = nn.Linear(input_dim, model_dim) | |
| encoder_layer = nn.TransformerEncoderLayer( | |
| d_model=model_dim, | |
| nhead=num_heads, | |
| dropout=dropout, | |
| dim_feedforward=2*model_dim, # 128 | |
| batch_first=True | |
| ) | |
| encoder_norm = nn.LayerNorm(model_dim) | |
| self.transformer_encoder = nn.TransformerEncoder( | |
| encoder_layer, | |
| num_layers=num_layers, | |
| norm=encoder_norm | |
| ) | |
| decoder_layer = nn.TransformerDecoderLayer( | |
| d_model=model_dim, | |
| nhead=num_heads, | |
| dropout=dropout, | |
| dim_feedforward=2*model_dim, # 128 | |
| batch_first=True | |
| ) | |
| decoder_norm = nn.LayerNorm(model_dim) | |
| self.transformer_decoder = nn.TransformerDecoder( | |
| decoder_layer, | |
| num_layers=num_layers, | |
| norm=decoder_norm | |
| ) | |
| self.output = nn.Linear(model_dim, input_dim) | |
| def forward(self, x): | |
| embed_x = self.embedding(x) | |
| enc_out = self.transformer_encoder(embed_x) | |
| dec_out = self.transformer_decoder(embed_x, enc_out) | |
| out = self.output(dec_out) | |
| return out | |