Spaces:
Runtime error
Runtime error
File size: 9,203 Bytes
1ccdd5a e0f96a9 1ccdd5a e0f96a9 1ccdd5a e0f96a9 1ccdd5a e0f96a9 1ccdd5a e0f96a9 1ccdd5a e0f96a9 1ccdd5a e0f96a9 1ccdd5a e0f96a9 1ccdd5a e0f96a9 1ccdd5a e0f96a9 1ccdd5a e0f96a9 1ccdd5a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 |
import math
import torch
import torch.nn as nn
from utils import normalize_data
import torch.nn.functional as F
from torch.nn import TransformerEncoder, TransformerEncoderLayer
class StyleEncoder(nn.Module):
def __init__(self, num_hyperparameters, em_size):
super().__init__()
self.em_size = em_size
self.embedding = nn.Linear(num_hyperparameters, self.em_size)
def forward(self, hyperparameters): # B x num_hps
return self.embedding(hyperparameters)
class StyleEmbEncoder(nn.Module):
def __init__(self, num_hyperparameters, em_size, num_embeddings=100):
super().__init__()
assert num_hyperparameters == 1
self.em_size = em_size
self.embedding = nn.Embedding(num_embeddings, self.em_size)
def forward(self, hyperparameters): # B x num_hps
return self.embedding(hyperparameters.squeeze(1))
class _PositionalEncoding(nn.Module):
def __init__(self, d_model, dropout=0.):
super().__init__()
self.dropout = nn.Dropout(p=dropout)
self.d_model = d_model
self.device_test_tensor = nn.Parameter(torch.tensor(1.))
def forward(self, x):# T x B x num_features
assert self.d_model % x.shape[-1]*2 == 0
d_per_feature = self.d_model // x.shape[-1]
pe = torch.zeros(*x.shape, d_per_feature, device=self.device_test_tensor.device)
#position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
interval_size = 10
div_term = (1./interval_size) * 2*math.pi*torch.exp(torch.arange(0, d_per_feature, 2, device=self.device_test_tensor.device).float()*math.log(math.sqrt(2)))
#print(div_term/2/math.pi)
pe[..., 0::2] = torch.sin(x.unsqueeze(-1) * div_term)
pe[..., 1::2] = torch.cos(x.unsqueeze(-1) * div_term)
return self.dropout(pe).view(x.shape[0],x.shape[1],self.d_model)
Positional = lambda _, emsize: _PositionalEncoding(d_model=emsize)
class EmbeddingEncoder(nn.Module):
def __init__(self, num_features, em_size, num_embs=100):
super().__init__()
self.num_embs = num_embs
self.embeddings = nn.Embedding(num_embs * num_features, em_size, max_norm=True)
self.init_weights(.1)
self.min_max = (-2,+2)
@property
def width(self):
return self.min_max[1] - self.min_max[0]
def init_weights(self, initrange):
self.embeddings.weight.data.uniform_(-initrange, initrange)
def discretize(self, x):
split_size = self.width / self.num_embs
return (x - self.min_max[0] // split_size).int().clamp(0, self.num_embs - 1)
def forward(self, x): # T x B x num_features
x_idxs = self.discretize(x)
x_idxs += torch.arange(x.shape[-1], device=x.device).view(1, 1, -1) * self.num_embs
# print(x_idxs,self.embeddings.weight.shape)
return self.embeddings(x_idxs).mean(-2)
class Normalize(nn.Module):
def __init__(self, mean, std):
super().__init__()
self.mean = mean
self.std = std
def forward(self, x):
return (x-self.mean)/self.std
def get_normalized_uniform_encoder(encoder_creator):
"""
This can be used to wrap an encoder that is fed uniform samples in [0,1] and normalizes these to 0 mean and 1 std.
For example, it can be used as `encoder_creator = get_normalized_uniform_encoder(encoders.Linear)`, now this can
be initialized with `encoder_creator(feature_dim, in_dim)`.
:param encoder:
:return:
"""
return lambda in_dim, out_dim: nn.Sequential(Normalize(.5, math.sqrt(1/12)), encoder_creator(in_dim, out_dim))
def get_normalized_encoder(encoder_creator, data_std):
return lambda in_dim, out_dim: nn.Sequential(Normalize(0., data_std), encoder_creator(in_dim, out_dim))
class ZNormalize(nn.Module):
def forward(self, x):
return (x-x.mean(-1,keepdim=True))/x.std(-1,keepdim=True)
class AppendEmbeddingEncoder(nn.Module):
def __init__(self, base_encoder, num_features, emsize):
super().__init__()
self.num_features = num_features
self.base_encoder = base_encoder
self.emb = nn.Parameter(torch.zeros(emsize))
def forward(self, x):
if (x[-1] == 1.).all():
append_embedding = True
else:
assert (x[-1] == 0.).all(), "You need to specify as last position whether to append embedding. " \
"If you don't want this behavior, please use the wrapped encoder instead."
append_embedding = False
x = x[:-1]
encoded_x = self.base_encoder(x)
if append_embedding:
encoded_x = torch.cat([encoded_x, self.emb[None, None, :].repeat(1, encoded_x.shape[1], 1)], 0)
return encoded_x
def get_append_embedding_encoder(encoder_creator):
return lambda num_features, emsize: AppendEmbeddingEncoder(encoder_creator(num_features, emsize), num_features, emsize)
class VariableNumFeaturesEncoder(nn.Module):
def __init__(self, base_encoder, num_features):
super().__init__()
self.base_encoder = base_encoder
self.num_features = num_features
def forward(self, x):
x = x * (self.num_features/x.shape[-1])
x = torch.cat((x, torch.zeros(*x.shape[:-1], self.num_features - x.shape[-1], device=x.device)), -1)
return self.base_encoder(x)
def get_variable_num_features_encoder(encoder_creator):
return lambda num_features, emsize: VariableNumFeaturesEncoder(encoder_creator(num_features, emsize), num_features)
class NoMeanEncoder(nn.Module):
"""
This can be useful for any prior that is translation invariant in x or y.
A standard GP for example is translation invariant in x.
That is, GP(x_test+const,x_train+const,y_train) = GP(x_test,x_train,y_train).
"""
def __init__(self, base_encoder):
super().__init__()
self.base_encoder = base_encoder
def forward(self, x):
return self.base_encoder(x - x.mean(0, keepdim=True))
def get_no_mean_encoder(encoder_creator):
return lambda num_features, emsize: NoMeanEncoder(encoder_creator(num_features, emsize))
Linear = nn.Linear
MLP = lambda num_features, emsize: nn.Sequential(nn.Linear(num_features+1,emsize*2),
nn.ReLU(),
nn.Linear(emsize*2,emsize))
class NanHandlingEncoder(nn.Module):
def __init__(self, num_features, emsize, keep_nans=True):
super().__init__()
self.num_features = 2 * num_features if keep_nans else num_features
self.emsize = emsize
self.keep_nans = keep_nans
self.layer = nn.Linear(self.num_features, self.emsize)
def forward(self, x):
if self.keep_nans:
x = torch.cat([torch.nan_to_num(x, nan=0.0), normalize_data(torch.isnan(x) * -1
+ torch.logical_and(torch.isinf(x), torch.sign(x) == 1) * 1
+ torch.logical_and(torch.isinf(x), torch.sign(x) == -1) * 2
)], -1)
else:
x = torch.nan_to_num(x, nan=0.0)
return self.layer(x)
class Linear(nn.Linear):
def __init__(self, num_features, emsize, replace_nan_by_zero=False):
super().__init__(num_features, emsize)
self.num_features = num_features
self.emsize = emsize
self.replace_nan_by_zero = replace_nan_by_zero
def forward(self, x):
if self.replace_nan_by_zero:
x = torch.nan_to_num(x, nan=0.0)
return super().forward(x)
def __setstate__(self, state):
super().__setstate__(state)
self.__dict__.setdefault('replace_nan_by_zero', True)
class Conv(nn.Module):
def __init__(self, input_size, emsize):
super().__init__()
self.convs = torch.nn.ModuleList([nn.Conv2d(64 if i else 1, 64, 3) for i in range(5)])
self.linear = nn.Linear(64,emsize)
def forward(self, x):
size = math.isqrt(x.shape[-1])
assert size*size == x.shape[-1]
x = x.reshape(*x.shape[:-1], 1, size, size)
for conv in self.convs:
if x.shape[-1] < 4:
break
x = conv(x)
x.relu_()
x = nn.AdaptiveAvgPool2d((1,1))(x).squeeze(-1).squeeze(-1)
return self.linear(x)
class CanEmb(nn.Embedding):
def __init__(self, num_features, num_embeddings: int, embedding_dim: int, *args, **kwargs):
assert embedding_dim % num_features == 0
embedding_dim = embedding_dim // num_features
super().__init__(num_embeddings, embedding_dim, *args, **kwargs)
def forward(self, x):
lx = x.long()
assert (lx == x).all(), "CanEmb only works with tensors of whole numbers"
x = super().forward(lx)
return x.view(*x.shape[:-2], -1)
def get_Canonical(num_classes):
return lambda num_features, emsize: CanEmb(num_features, num_classes, emsize)
def get_Embedding(num_embs_per_feature=100):
return lambda num_features, emsize: EmbeddingEncoder(num_features, emsize, num_embs=num_embs_per_feature)
|