|
| 1 | +import numpy as np |
| 2 | +import glob |
| 3 | +import torch |
| 4 | +import torch.nn as nn |
| 5 | +from sklearn import neighbors |
| 6 | +import logging |
| 7 | + |
| 8 | + |
| 9 | +class PositionalEncoding(nn.Module): |
| 10 | + def __init__(self, d_emb, norm=10000.0): |
| 11 | + """ |
| 12 | + Inputs |
| 13 | + d_model - Hidden dimensionality. |
| 14 | + """ |
| 15 | + super().__init__() |
| 16 | + self.d_emb = d_emb |
| 17 | + self.norm = norm |
| 18 | + |
| 19 | + def forward(self, t): |
| 20 | + pe = torch.zeros(t.shape[0], t.shape[1], self.d_emb).to(t.device) # (B, T, D) |
| 21 | + div_term = torch.exp(torch.arange(0, self.d_emb, 2).float() * (-np.log(self.norm) / self.d_emb))[None, None, :].to(t.device) # (1, 1, D / 2) |
| 22 | + t = t.unsqueeze(2) # (B, 1, T) |
| 23 | + pe[:, :, 0::2] = torch.sin(t * div_term) # (B, T, D / 2) |
| 24 | + pe[:, :, 1::2] = torch.cos(t * div_term) # (B, T, D / 2) |
| 25 | + return pe # (B, T, D) |
| 26 | + |
| 27 | +class TransformerModel(nn.Module): |
| 28 | + |
| 29 | + def __init__(self, ninp, nemb, nout, nhead, nhid, nlayers, dropout=0.1, norm=1000.0): |
| 30 | + """ |
| 31 | + Transformer model for sequence to sequence learning |
| 32 | +
|
| 33 | + Args: |
| 34 | + ninp (_type_): input size |
| 35 | + nemb (_type_): embedding size |
| 36 | + nout (_type_): output size |
| 37 | + nhead (_type_): number of heads |
| 38 | + nhid (_type_): hidden layer size in feed forward network |
| 39 | + nlayers (_type_): number of layers |
| 40 | + dropout (float, optional): dropout probability. Defaults to 0.5. |
| 41 | + """ |
| 42 | + super(TransformerModel, self).__init__() |
| 43 | + |
| 44 | + self.model_type = 'Transformer' |
| 45 | + |
| 46 | + self.encoder = nn.Linear(ninp, nemb) |
| 47 | + |
| 48 | + self.pos_encoder = PositionalEncoding(nemb, norm) |
| 49 | + |
| 50 | + encoder_layers = nn.TransformerEncoderLayer(nemb, nhead, nhid, dropout, norm_first=True, batch_first=True) |
| 51 | + self.transformer_encoder = nn.TransformerEncoder(encoder_layers, nlayers, enable_nested_tensor=False) |
| 52 | + |
| 53 | + self.nemb = nemb |
| 54 | + self.decoder = nn.Linear(nemb, nout) |
| 55 | + |
| 56 | + self.init_weights() |
| 57 | + |
| 58 | + def init_weights(self): |
| 59 | + |
| 60 | + # Since TraansformerEncoder inputs a TransformerEncoderLayer, all layers will use exactly the same initialization |
| 61 | + # We undo this here |
| 62 | + for name, param in self.named_parameters(): |
| 63 | + if 'weight' in name and param.data.dim() == 2: |
| 64 | + nn.init.kaiming_uniform_(param) |
| 65 | + |
| 66 | + def forward(self, src, tau, src_mask): |
| 67 | + |
| 68 | + # Get tau embedding |
| 69 | + tau_emb = self.pos_encoder(tau) |
| 70 | + |
| 71 | + # Embed the input sequence into the embedding space and add the tau embedding |
| 72 | + x = self.encoder(src) + tau_emb |
| 73 | + |
| 74 | + # Apply the transformer encoder |
| 75 | + x = self.transformer_encoder(x, src_key_padding_mask=src_mask) |
| 76 | + |
| 77 | + # Apply the decoder to the output space |
| 78 | + x = self.decoder(x) |
| 79 | + |
| 80 | + output = (~src_mask).float()[:, :, None] * x |
| 81 | + |
| 82 | + return output |
| 83 | + |
| 84 | + |
| 85 | +class Forward(object): |
| 86 | + def __init__(self, gpu=0, checkpoint=None, readir=None, verbose=0): |
| 87 | + |
| 88 | + self.logger = logging.getLogger("neural") |
| 89 | + self.logger.setLevel(logging.DEBUG) |
| 90 | + self.logger.handlers = [] |
| 91 | + ch = logging.StreamHandler() |
| 92 | + formatter = logging.Formatter('%(asctime)s - %(message)s') |
| 93 | + ch.setFormatter(formatter) |
| 94 | + self.logger.addHandler(ch) |
| 95 | + |
| 96 | + # Is a GPU available? |
| 97 | + self.cuda = torch.cuda.is_available() |
| 98 | + self.gpu = gpu |
| 99 | + self.device = torch.device("cpu") #f"cuda:{self.gpu}" if self.cuda else "cpu") |
| 100 | + |
| 101 | + if (checkpoint is None): |
| 102 | + if readir is None: |
| 103 | + raise ValueError('Not checkpoint or read directory selected') |
| 104 | + files = glob.glob(readir + '*.pth') |
| 105 | + self.checkpoint = sorted(files)[-1] |
| 106 | + else: |
| 107 | + self.checkpoint = checkpoint |
| 108 | + |
| 109 | + checkpoint = torch.load(self.checkpoint, map_location=lambda storage, loc: storage, weights_only=False) |
| 110 | + |
| 111 | + self.hyperparameters = checkpoint['hyperparameters'] |
| 112 | + self.predict_model = TransformerModel(ninp=self.hyperparameters['transformer']['n_input'], |
| 113 | + nemb=self.hyperparameters['transformer']['n_embedding'], |
| 114 | + nout=self.hyperparameters['transformer']['n_output'], |
| 115 | + nhead=self.hyperparameters['transformer']['n_heads'], |
| 116 | + nhid=self.hyperparameters['transformer']['n_hidden'], |
| 117 | + nlayers=self.hyperparameters['transformer']['n_layers'], |
| 118 | + norm=self.hyperparameters['transformer']['norm'], |
| 119 | + dropout=self.hyperparameters['transformer']['dropout']).to(self.device) |
| 120 | + self.predict_model.load_state_dict(checkpoint['state_dict']) |
| 121 | + |
| 122 | + self.predict_model.eval() |
| 123 | + |
| 124 | + if (verbose >= 1): |
| 125 | + npars = sum(p.numel() for p in self.predict_model.parameters() if p.requires_grad) |
| 126 | + tmp = self.checkpoint.split('/') |
| 127 | + self.logger.info(f' * Using neural checkpoint {tmp[-1]} on {self.device} - N. parameters = {npars}') |
| 128 | + |
| 129 | + def predict(self, tau_all, ne_all, vturb_all, T_all, vlos_all): |
| 130 | + |
| 131 | + tau = (np.log10(tau_all.astype('float32')) + 10.0) * 10.0 |
| 132 | + vturb = vturb_all.astype('float32') / 1e3 - 6.0 |
| 133 | + vlos = vlos_all.astype('float32') / 1e3 |
| 134 | + T = np.log10(T_all.astype('float32')) - 3.8 |
| 135 | + ne = np.log10(ne_all.astype('float32')) - 16.0 |
| 136 | + |
| 137 | + pars = np.concatenate([vturb[None, :], vlos[None, :], T[None, :], ne[None, :]], axis=0).T |
| 138 | + mask = np.zeros(len(tau)).astype('bool') |
| 139 | + |
| 140 | + pars = torch.tensor(pars, dtype=torch.float32).to(self.device) |
| 141 | + tau = torch.tensor(tau, dtype=torch.float32).to(self.device) |
| 142 | + mask = torch.tensor(mask, dtype=torch.bool).to(self.device) |
| 143 | + |
| 144 | + with torch.no_grad(): |
| 145 | + self.pred_out = self.predict_model(pars[None, ...], tau[None, ...], mask[None, ...]) |
| 146 | + |
| 147 | + return self.pred_out[0, ...].cpu().numpy() |
0 commit comments