engression 0.0.0.dev0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- engression/__init__.py +10 -0
- engression/data/__init__.py +0 -0
- engression/data/loader.py +51 -0
- engression/data/simulator.py +55 -0
- engression/engression.py +358 -0
- engression/engression_bagged.py +282 -0
- engression/loss_func.py +52 -0
- engression/models.py +179 -0
- engression/utils.py +65 -0
- engression-0.0.0.dev0.dist-info/METADATA +47 -0
- engression-0.0.0.dev0.dist-info/RECORD +13 -0
- engression-0.0.0.dev0.dist-info/WHEEL +5 -0
- engression-0.0.0.dev0.dist-info/top_level.txt +1 -0
engression/__init__.py
ADDED
|
@@ -0,0 +1,10 @@
|
|
|
1
|
+
from .engression import engression
|
|
2
|
+
# from .engression_bagged import engression_bagged
|
|
3
|
+
|
|
4
|
+
try:
|
|
5
|
+
# pylint: disable=wrong-import-position
|
|
6
|
+
import torch
|
|
7
|
+
except ModuleNotFoundError:
|
|
8
|
+
raise ModuleNotFoundError(
|
|
9
|
+
"No module named 'torch', and engression depends on PyTorch (aka 'torch')."
|
|
10
|
+
"Visit https://pytorch.org/ for installation instructions.")
|
|
File without changes
|
|
@@ -0,0 +1,51 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
from torch.utils.data import TensorDataset, DataLoader
|
|
3
|
+
|
|
4
|
+
def make_dataloader(x, y, batch_size=128, shuffle=True, num_workers=0):
|
|
5
|
+
"""Make dataloader.
|
|
6
|
+
|
|
7
|
+
Args:
|
|
8
|
+
x (torch.Tensor): data of predictors.
|
|
9
|
+
y (torch.Tensor): data of responses.
|
|
10
|
+
batch_size (int, optional): _description_. Defaults to 128.
|
|
11
|
+
shuffle (bool, optional): _description_. Defaults to True.
|
|
12
|
+
num_workers (int, optional): _description_. Defaults to 0.
|
|
13
|
+
|
|
14
|
+
Returns:
|
|
15
|
+
_type_: _description_
|
|
16
|
+
"""
|
|
17
|
+
dataset = TensorDataset(x, y)
|
|
18
|
+
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)
|
|
19
|
+
return dataloader
|
|
20
|
+
|
|
21
|
+
def partition_data(x_full, y_full, cut_quantile=0.3, split_train="smaller"):
|
|
22
|
+
"""Partition data into training and test sets.
|
|
23
|
+
|
|
24
|
+
Args:
|
|
25
|
+
x_full (torch.Tensor): full data of x.
|
|
26
|
+
y_full (torch.Tensor): full data of y.
|
|
27
|
+
cut_quantile (float, optional): quantile of the cutting point of x. Defaults to 0.3.
|
|
28
|
+
split_train (str, optional): which subset is used for for training. choices=["smaller", "larger"]. Defaults to "smaller".
|
|
29
|
+
|
|
30
|
+
Returns:
|
|
31
|
+
tuple of torch.Tensors: training and test data.
|
|
32
|
+
"""
|
|
33
|
+
# Split data into training and test sets.
|
|
34
|
+
x_cut = torch.quantile(x_full, cut_quantile)
|
|
35
|
+
train_idx = x_full <= x_cut if split_train == "smaller" else x_full >= x_cut
|
|
36
|
+
x_tr = x_full[train_idx]
|
|
37
|
+
y_tr = y_full[train_idx]
|
|
38
|
+
x_te = x_full[~train_idx]
|
|
39
|
+
y_te = y_full[~train_idx]
|
|
40
|
+
|
|
41
|
+
# Standardize data based on training statistics.
|
|
42
|
+
x_tr_mean = x_tr.mean()
|
|
43
|
+
x_tr_std = x_tr.std()
|
|
44
|
+
y_tr_mean = y_tr.mean()
|
|
45
|
+
y_tr_std = y_tr.std()
|
|
46
|
+
x_tr = (x_tr - x_tr_mean)/x_tr_std
|
|
47
|
+
y_tr = (y_tr - y_tr_mean)/y_tr_std
|
|
48
|
+
x_te = (x_te - x_tr_mean)/x_tr_std
|
|
49
|
+
y_te = (y_te - y_tr_mean)/y_tr_std
|
|
50
|
+
x_full_normal = (x_full - x_tr_mean)/x_tr_std
|
|
51
|
+
return x_tr.unsqueeze(1), y_tr.unsqueeze(1), x_te.unsqueeze(1), y_te.unsqueeze(1), x_full_normal
|
|
@@ -0,0 +1,55 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
import torch.nn as nn
|
|
3
|
+
import numpy as np
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
def preanm_simulator(true_function='softplus', n=10000, x_lower=0, x_upper=2, noise_std=1, noise_dist='gaussian', train=True, device=torch.device('cpu')):
|
|
7
|
+
"""Data simulator for a pre-additive noise model (pre-ANM).
|
|
8
|
+
|
|
9
|
+
Args:
|
|
10
|
+
true_function (str, optional): true function g^\star. Defaults to 'softplus'. Choices: ['softplus', 'cubic','square', 'log'].
|
|
11
|
+
n (int, optional): sample size. Defaults to 10000.
|
|
12
|
+
x_lower (int, optional): lower bound of the training support. Defaults to 0.
|
|
13
|
+
x_upper (int, optional): upper bound of the training support. Defaults to 2.
|
|
14
|
+
noise_std (int, optional): standard deviation of the noise. Defaults to 1.
|
|
15
|
+
noise_dist (str, optional): noise distribution. Defaults to 'gaussian'. Choices: ['gaussian', 'uniform'].
|
|
16
|
+
train (bool, optional): generate data for training. Defaults to True.
|
|
17
|
+
device (str or torch.device, optional): _description_. Defaults to torch.device('cpu').
|
|
18
|
+
|
|
19
|
+
Returns:
|
|
20
|
+
tuple of torch.Tensors: data simulated from a pre-ANM.
|
|
21
|
+
"""
|
|
22
|
+
if isinstance(true_function, str):
|
|
23
|
+
if true_function == 'softplus':
|
|
24
|
+
true_function = lambda x: nn.Softplus()(x)
|
|
25
|
+
elif true_function == 'cubic':
|
|
26
|
+
true_function = lambda x: x.pow(3)/3
|
|
27
|
+
elif true_function == 'square':
|
|
28
|
+
true_function = lambda x: (nn.functional.relu(x)).pow(2)/2
|
|
29
|
+
elif true_function == 'log':
|
|
30
|
+
true_function = lambda x: (x/3 + np.log(3) - 2/3)*(x <= 2) + (torch.log(1 + x*(x > 2)))*(x > 2)
|
|
31
|
+
|
|
32
|
+
if isinstance(device, str):
|
|
33
|
+
device = torch.device(device)
|
|
34
|
+
|
|
35
|
+
if train:
|
|
36
|
+
x = torch.rand(n, 1)*(x_upper - x_lower) + x_lower
|
|
37
|
+
if noise_dist == 'gaussian':
|
|
38
|
+
eps = torch.randn(n, 1)*noise_std
|
|
39
|
+
else:
|
|
40
|
+
assert noise_dist == 'uniform'
|
|
41
|
+
eps = (torch.rand(n, 1) - 0.5)*noise_std*np.sqrt(12)
|
|
42
|
+
xn = x + eps
|
|
43
|
+
y = true_function(xn)
|
|
44
|
+
return x.to(device), y.to(device)
|
|
45
|
+
|
|
46
|
+
else:
|
|
47
|
+
x_eval = torch.linspace(x_lower, x_upper, n).unsqueeze(1)
|
|
48
|
+
y_eval_med = true_function(x_eval)
|
|
49
|
+
gen_sample_size = 10000
|
|
50
|
+
x_rep = torch.repeat_interleave(x_eval, (gen_sample_size * torch.ones(n)).long(), dim=0)
|
|
51
|
+
x_rep = x_rep + torch.randn(x_rep.size(0), 1)*noise_std
|
|
52
|
+
y_eval_mean = true_function(x_rep)
|
|
53
|
+
y_eval_mean = list(torch.split(y_eval_mean, gen_sample_size))
|
|
54
|
+
y_eval_mean = torch.cat([y_eval_mean[i].mean().unsqueeze(0) for i in range(n)], dim=0).unsqueeze(1)
|
|
55
|
+
return x_eval.to(device), y_eval_med.to(device), y_eval_mean.to(device)
|
engression/engression.py
ADDED
|
@@ -0,0 +1,358 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
import matplotlib.pyplot as plt
|
|
3
|
+
|
|
4
|
+
from .loss_func import *
|
|
5
|
+
from .models import StoNet
|
|
6
|
+
from .data.loader import make_dataloader
|
|
7
|
+
from .utils import *
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class Engressor(object):
|
|
11
|
+
|
|
12
|
+
def __init__(self,
|
|
13
|
+
in_dim, out_dim, num_layer=2, hidden_dim=100, noise_dim=100,
|
|
14
|
+
lr=0.001, num_epoches=500, batch_size=None, device="cpu", standardize=True):
|
|
15
|
+
"""Engressor class
|
|
16
|
+
|
|
17
|
+
Args:
|
|
18
|
+
in_dim (int): input dimension
|
|
19
|
+
out_dim (int): output dimension
|
|
20
|
+
num_layer (int, optional): number of layers. Defaults to 2.
|
|
21
|
+
hidden_dim (int, optional): number of neurons per layer. Defaults to 100.
|
|
22
|
+
noise_dim (int, optional): noise dimension. Defaults to 100.
|
|
23
|
+
lr (float, optional): learning rate. Defaults to 0.001.
|
|
24
|
+
num_epoches (int, optional): number of epoches. Defaults to 500.
|
|
25
|
+
batch_size (int, optional): batch size. Defaults to None, referring to the full batch.
|
|
26
|
+
device (str or torch.device, optional): device. Defaults to "cpu". Choices = ["cpu", "cuda"].
|
|
27
|
+
standardize (bool, optional): whether to standardize data. Defaults to True.
|
|
28
|
+
"""
|
|
29
|
+
super().__init__()
|
|
30
|
+
self.num_layer = num_layer
|
|
31
|
+
self.hidden_dim = hidden_dim
|
|
32
|
+
self.noise_dim = noise_dim
|
|
33
|
+
self.lr = lr
|
|
34
|
+
self.num_epoches = num_epoches
|
|
35
|
+
self.batch_size = batch_size
|
|
36
|
+
if isinstance(device, str):
|
|
37
|
+
if device == "gpu" or device == "cuda":
|
|
38
|
+
device = torch.device("cuda")
|
|
39
|
+
else:
|
|
40
|
+
device = torch.device(device)
|
|
41
|
+
self.device = device
|
|
42
|
+
check_for_gpu(self.device)
|
|
43
|
+
self.standardize = standardize
|
|
44
|
+
self.x_mean = None
|
|
45
|
+
self.x_std = None
|
|
46
|
+
self.y_mean = None
|
|
47
|
+
self.y_std = None
|
|
48
|
+
|
|
49
|
+
self.model = StoNet(in_dim, out_dim, num_layer, hidden_dim, noise_dim).to(self.device)
|
|
50
|
+
self.optimizer = torch.optim.Adam(self.model.parameters(), lr=lr)
|
|
51
|
+
|
|
52
|
+
self.tr_loss = None
|
|
53
|
+
|
|
54
|
+
def train_mode(self):
|
|
55
|
+
self.model.train()
|
|
56
|
+
|
|
57
|
+
def eval_mode(self):
|
|
58
|
+
self.model.eval()
|
|
59
|
+
|
|
60
|
+
def summary(self):
|
|
61
|
+
"""Print the model architecture and hyperparameters."""
|
|
62
|
+
print("Engression model with\n" +
|
|
63
|
+
"\t number of layers: {}\n".format(self.num_layer) +
|
|
64
|
+
"\t hidden dimensions: {}\n".format(self.hidden_dim) +
|
|
65
|
+
"\t noise dimensions: {}\n".format(self.noise_dim) +
|
|
66
|
+
"\t number of epochs: {}\n".format(self.num_epoches) +
|
|
67
|
+
"\t batch size: {}\n".format(self.batch_size) +
|
|
68
|
+
"\t learning rate: {}\n".format(self.lr) +
|
|
69
|
+
"\t standardization: {}\n".format(self.standardize) +
|
|
70
|
+
"\t training mode: {}\n".format(self.model.training) +
|
|
71
|
+
"\t device: {}\n".format(self.device))
|
|
72
|
+
print("Training loss (original scale):\n" +
|
|
73
|
+
"\t energy-loss: {:.2f}, \n\tE(|Y-Yhat|): {:.2f}, \n\tE(|Yhat-Yhat'|): {:.2f}".format(
|
|
74
|
+
self.tr_loss[0], self.tr_loss[1], self.tr_loss[2]))
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
def _standardize_data_and_record_stats(self, x, y):
|
|
78
|
+
"""Standardize the data and record the mean and standard deviation of the training data.
|
|
79
|
+
|
|
80
|
+
Args:
|
|
81
|
+
x (torch.Tensor): training data of predictors.
|
|
82
|
+
y (torch.Tensor): training data of responses.
|
|
83
|
+
|
|
84
|
+
Returns:
|
|
85
|
+
torch.Tensor: standardized data.
|
|
86
|
+
"""
|
|
87
|
+
self.x_mean = torch.mean(x, dim=0)
|
|
88
|
+
self.x_std = torch.std(x, dim=0)
|
|
89
|
+
self.y_mean = torch.mean(y, dim=0)
|
|
90
|
+
self.y_std = torch.std(y, dim=0)
|
|
91
|
+
return (x - self.x_mean) / self.x_std, (y - self.y_mean) / self.y_std
|
|
92
|
+
|
|
93
|
+
def standardize_data(self, x, y=None):
|
|
94
|
+
"""Standardize the data, if self.standardize is True.
|
|
95
|
+
|
|
96
|
+
Args:
|
|
97
|
+
x (torch.Tensor): training data of predictors.
|
|
98
|
+
y (torch.Tensor, optional): _description_. Defaults to None.
|
|
99
|
+
|
|
100
|
+
Returns:
|
|
101
|
+
torch.Tensor: standardized or original data.
|
|
102
|
+
"""
|
|
103
|
+
if y is None:
|
|
104
|
+
if self.standardize:
|
|
105
|
+
return (x - self.x_mean) / self.x_std
|
|
106
|
+
else:
|
|
107
|
+
return x
|
|
108
|
+
else:
|
|
109
|
+
if self.standardize:
|
|
110
|
+
return (x - self.x_mean) / self.x_std, (y - self.y_mean) / self.y_std
|
|
111
|
+
else:
|
|
112
|
+
return x, y
|
|
113
|
+
|
|
114
|
+
def unstandardize_data(self, y, x=None):
|
|
115
|
+
"""Transform the predictions back to the original scale, if self.standardize is True.
|
|
116
|
+
|
|
117
|
+
Args:
|
|
118
|
+
y (torch.Tensor): data in the standardized scale
|
|
119
|
+
|
|
120
|
+
Returns:
|
|
121
|
+
torch.Tensor: data in the original scale
|
|
122
|
+
"""
|
|
123
|
+
if x is None:
|
|
124
|
+
if self.standardize:
|
|
125
|
+
return y * self.y_std + self.y_mean
|
|
126
|
+
else:
|
|
127
|
+
return y
|
|
128
|
+
else:
|
|
129
|
+
if self.standardize:
|
|
130
|
+
return x * self.x_std + self.x_mean, y * self.y_std + self.y_mean
|
|
131
|
+
else:
|
|
132
|
+
return x, y
|
|
133
|
+
|
|
134
|
+
def train(self, x, y, num_epoches=None, batch_size=512, print_every_nepoch=100, print_times_per_epoch=1, standardize=True, verbose=True):
|
|
135
|
+
"""Training function.
|
|
136
|
+
|
|
137
|
+
Args:
|
|
138
|
+
x (torch.Tensor): training data of predictors.
|
|
139
|
+
y (torch.Tensor): trainging data of responses.
|
|
140
|
+
num_epoches (int, optional): number of training epochs. Defaults to None.
|
|
141
|
+
batch_size (int, optional): batch size for mini-batch SGD. Defaults to 512.
|
|
142
|
+
print_every_nepoch (int, optional): print losses every print_every_nepoch number of epochs. Defaults to 100.
|
|
143
|
+
print_times_per_epoch (int, optional): print losses for print_times_per_epoch times per epoch. Defaults to 1.
|
|
144
|
+
standardize (bool, optional): standardize the data. Defaults to True.
|
|
145
|
+
"""
|
|
146
|
+
self.train_mode()
|
|
147
|
+
if num_epoches is None:
|
|
148
|
+
num_epoches = self.num_epoches
|
|
149
|
+
if batch_size is None:
|
|
150
|
+
batch_size = self.batch_size
|
|
151
|
+
if standardize:
|
|
152
|
+
self.standardize = standardize
|
|
153
|
+
|
|
154
|
+
x = vectorize(x)
|
|
155
|
+
y = vectorize(y)
|
|
156
|
+
x = x.to(self.device)
|
|
157
|
+
y = y.to(self.device)
|
|
158
|
+
if self.standardize:
|
|
159
|
+
print("Data is standardized during training only; the printed training losses are on the standardized scale, \n" +
|
|
160
|
+
"while during evaluation, the predictions, evaluation metrics, and plots will be on the original scale.\n")
|
|
161
|
+
x, y = self._standardize_data_and_record_stats(x, y)
|
|
162
|
+
|
|
163
|
+
if batch_size >= x.size(0)//2:
|
|
164
|
+
print("Batch is larger than half of the sample size. Training based on full-batch gradient descent.")
|
|
165
|
+
self.batch_size = x.size(0)
|
|
166
|
+
for epoch_idx in range(num_epoches):
|
|
167
|
+
self.model.zero_grad()
|
|
168
|
+
y_sample1 = self.model(x)
|
|
169
|
+
y_sample2 = self.model(x)
|
|
170
|
+
loss, loss1, loss2 = energy_loss_two_sample(y, y_sample1, y_sample2, verbose=verbose)
|
|
171
|
+
loss.backward()
|
|
172
|
+
self.optimizer.step()
|
|
173
|
+
if epoch_idx == 0 or (epoch_idx + 1) % print_every_nepoch == 0:
|
|
174
|
+
print("[Epoch {} ({:.0f}%)] energy-loss: {:.4f}, E(|Y-Yhat|): {:.4f}, E(|Yhat-Yhat'|): {:.4f}".format(
|
|
175
|
+
epoch_idx + 1, 100 * epoch_idx / num_epoches, loss.item(), loss1.item(), loss2.item()))
|
|
176
|
+
else:
|
|
177
|
+
train_loader = make_dataloader(x, y, batch_size=batch_size, shuffle=True)
|
|
178
|
+
print("Training based on mini-batch gradient descent with a batch size of {}.".format(batch_size))
|
|
179
|
+
for epoch_idx in range(num_epoches):
|
|
180
|
+
for batch_idx, (x_batch, y_batch) in enumerate(train_loader):
|
|
181
|
+
self.model.zero_grad()
|
|
182
|
+
y_sample1 = self.model(x_batch)
|
|
183
|
+
y_sample2 = self.model(x_batch)
|
|
184
|
+
loss, loss1, loss2 = energy_loss_two_sample(y_batch, y_sample1, y_sample2, verbose=verbose)
|
|
185
|
+
# loss = (y_batch - y_sample1).pow(2).mean()
|
|
186
|
+
loss.backward()
|
|
187
|
+
self.optimizer.step()
|
|
188
|
+
if epoch_idx == 0 or (epoch_idx + 1) % print_every_nepoch == 0:
|
|
189
|
+
if (batch_idx + 1) % (len(train_loader) // print_times_per_epoch) == 0:
|
|
190
|
+
print("[Epoch {} ({:.0f}%), batch {}]: energy-loss: {:.4f}, E(|Y-Yhat|): {:.4f}, E(|Yhat-Yhat'|): {:.4f}".format(
|
|
191
|
+
epoch_idx + 1, 100 * epoch_idx / num_epoches, batch_idx + 1, loss.item(), loss1.item(), loss2.item()))
|
|
192
|
+
# print("[Epoch {}, batch {}]: loss: {:.4f}".format(
|
|
193
|
+
# epoch_idx + 1, batch_idx + 1, loss.item()))
|
|
194
|
+
|
|
195
|
+
# Evaluate performance on the training data (on the original scale)
|
|
196
|
+
self.model.eval()
|
|
197
|
+
x, y = self.unstandardize_data(y, x)
|
|
198
|
+
self.tr_loss = self.eval_loss(x, y, loss_type="energy", verbose=True)
|
|
199
|
+
|
|
200
|
+
print("\nTraining loss on the original (non-standardized) scale:\n" +
|
|
201
|
+
"\tEnergy-loss: {:.4f}, E(|Y-Yhat|): {:.4f}, E(|Yhat-Yhat'|): {:.4f}".format(
|
|
202
|
+
self.tr_loss[0], self.tr_loss[1], self.tr_loss[2]))
|
|
203
|
+
|
|
204
|
+
print("\nPrediction-loss E(|Y-Yhat|) and variance-loss E(|Yhat-Yhat'|) should ideally be equally large" +
|
|
205
|
+
"\n-- consider training for more epochs or adjusting hyperparameters if there is a mismatch ")
|
|
206
|
+
|
|
207
|
+
def predict(self, x, target="mean", sample_size=100):
|
|
208
|
+
"""Point prediction.
|
|
209
|
+
|
|
210
|
+
Args:
|
|
211
|
+
x (torch.Tensor): data of predictors.
|
|
212
|
+
target (str or float or list, optional): single-valued functional to predict. float refers to the quantiles. Defaults to ["mean"].
|
|
213
|
+
sample_size (int, optional): sample sizes for each x. Defaults to 100.
|
|
214
|
+
|
|
215
|
+
Returns:
|
|
216
|
+
torch.Tensor or list of torch.Tensor: point predictions.
|
|
217
|
+
"""
|
|
218
|
+
self.eval_mode()
|
|
219
|
+
x = vectorize(x)
|
|
220
|
+
x = x.to(self.device)
|
|
221
|
+
x = self.standardize_data(x)
|
|
222
|
+
y_pred = self.model.predict(x, target, sample_size)
|
|
223
|
+
y_pred = self.unstandardize_data(y_pred)
|
|
224
|
+
return y_pred
|
|
225
|
+
|
|
226
|
+
def sample(self, x, sample_size=100, expand_dim=True):
|
|
227
|
+
"""Sample new response data.
|
|
228
|
+
|
|
229
|
+
Args:
|
|
230
|
+
x (torch.Tensor): test data of predictors.
|
|
231
|
+
target (str or float or list, optional): single-valued functional to predict. float refers to the quantiles. Defaults to ["mean"].
|
|
232
|
+
sample_size (int, optional): sample sizes for each x. Defaults to 100.
|
|
233
|
+
|
|
234
|
+
Returns:
|
|
235
|
+
torch.Tensor or list of torch.Tensor: samples.
|
|
236
|
+
- [:,:,i] gives the i-th sample of all x.
|
|
237
|
+
- [i,:,:] gives all samples of x_i.
|
|
238
|
+
"""
|
|
239
|
+
self.eval_mode()
|
|
240
|
+
x = vectorize(x)
|
|
241
|
+
x = x.to(self.device)
|
|
242
|
+
x = self.standardize_data(x)
|
|
243
|
+
y_samples = self.model.sample(x, sample_size, expand_dim=expand_dim)
|
|
244
|
+
y_samples = self.unstandardize_data(y_samples)
|
|
245
|
+
return y_samples
|
|
246
|
+
|
|
247
|
+
def eval_loss(self, x, y, loss_type="l2", sample_size=None, verbose=False):
|
|
248
|
+
"""Compute the loss for evaluation.
|
|
249
|
+
|
|
250
|
+
Args:
|
|
251
|
+
x (torch.Tensor): data of predictors.
|
|
252
|
+
y (torch.Tensor): data of responses.
|
|
253
|
+
loss_type (str, optional): loss type. Defaults to "l2". Choices: ["l2", "l1", "energy", "cor"].
|
|
254
|
+
sample_size (int, optional): sample sizes for each x. Defaults to 100.
|
|
255
|
+
|
|
256
|
+
Returns:
|
|
257
|
+
float: evaluation loss.
|
|
258
|
+
"""
|
|
259
|
+
if sample_size is None:
|
|
260
|
+
sample_size = 2 if loss_type == "energy" else 100
|
|
261
|
+
self.eval_mode()
|
|
262
|
+
x = vectorize(x)
|
|
263
|
+
y = vectorize(y)
|
|
264
|
+
x = x.to(self.device)
|
|
265
|
+
y = y.to(self.device)
|
|
266
|
+
if loss_type == "l2":
|
|
267
|
+
y_pred = self.predict(x, target="mean", sample_size=sample_size)
|
|
268
|
+
loss = (y - y_pred).pow(2).mean()
|
|
269
|
+
elif loss_type == "cor":
|
|
270
|
+
y_pred = self.predict(x, target="mean", sample_size=sample_size)
|
|
271
|
+
loss = cor(y, y_pred)
|
|
272
|
+
elif loss_type == "l1":
|
|
273
|
+
y_pred = self.predict(x, target=0.5, sample_size=sample_size)
|
|
274
|
+
loss = (y - y_pred).abs().mean()
|
|
275
|
+
else:
|
|
276
|
+
assert loss_type == "energy"
|
|
277
|
+
y_samples = self.sample(x, sample_size=sample_size, expand_dim=False)
|
|
278
|
+
loss = energy_loss(y, y_samples, verbose=verbose)
|
|
279
|
+
if not verbose:
|
|
280
|
+
return loss.item()
|
|
281
|
+
else:
|
|
282
|
+
loss, loss1, loss2 = loss
|
|
283
|
+
return loss.item(), loss1.item(), loss2.item()
|
|
284
|
+
|
|
285
|
+
def plot(self, x_te, y_te, x_tr=None, y_tr=None, x_idx=0, y_idx=0,
|
|
286
|
+
target="mean", sample_size=100, save_dir=None,
|
|
287
|
+
alpha=0.8, ymin=None, ymax=None):
|
|
288
|
+
"""Plot true data and predictions.
|
|
289
|
+
|
|
290
|
+
Args:
|
|
291
|
+
x_te (torch.Tensor): test data of predictors
|
|
292
|
+
y_te (torch.Tensor): test data of responses
|
|
293
|
+
x_tr (torch.Tensor): training data of predictors
|
|
294
|
+
y_tr (torch.Tensor): training data of responses
|
|
295
|
+
x_idx (int, optional): index of the predictor to plot (if there are multiple). Defaults to 0.
|
|
296
|
+
y_idx (int, optional): index of the response to plot (if there are multiple). Defaults to 0.
|
|
297
|
+
target (str or float, optional): target quantity. Defaults to "mean". Choice: ["mean", "median", "sample", float].
|
|
298
|
+
sample_size (int, optional): sample sizes for each x. Defaults to 100.
|
|
299
|
+
save_dir (str, optional): directory to save the plot. Defaults to None.
|
|
300
|
+
alpha (float, optional): transparency of the sampled data points. Defaults to 0.8.
|
|
301
|
+
ymin (float, optional): minimum value of y in the plot. Defaults to None.
|
|
302
|
+
ymax (float, optional): maximum value of y in the plot. Defaults to None.
|
|
303
|
+
"""
|
|
304
|
+
if x_tr is not None and y_tr is not None:
|
|
305
|
+
# Plot training data as well.
|
|
306
|
+
x_tr = vectorize(x_tr)
|
|
307
|
+
y_tr = vectorize(y_tr)
|
|
308
|
+
plt.scatter(x_tr[:,x_idx].cpu(), y_tr[:,y_idx].cpu(), s=1, label="training data", color="silver")
|
|
309
|
+
plt.scatter(x_te[:,x_idx].cpu(), y_te[:,y_idx].cpu(), s=1, label="test data", color="gold")
|
|
310
|
+
x = torch.cat((x_tr, x_te), dim=0)
|
|
311
|
+
y = torch.cat((y_tr, y_te), dim=0)
|
|
312
|
+
else:
|
|
313
|
+
# Plot only the test data.
|
|
314
|
+
x_te = vectorize(x_te)
|
|
315
|
+
y_te = vectorize(y_te)
|
|
316
|
+
plt.scatter(x_te[:,x_idx].cpu(), y_te[:,y_idx].cpu(), s=1, label="true data", color="silver")
|
|
317
|
+
x = x_te
|
|
318
|
+
y = y_te
|
|
319
|
+
x = x.to(self.device)
|
|
320
|
+
y = y.to(self.device)
|
|
321
|
+
|
|
322
|
+
if target != "sample":
|
|
323
|
+
y_pred = self.predict(x, target=target, sample_size=sample_size)
|
|
324
|
+
plt.scatter(x[:,x_idx].cpu(), y_pred[:,y_idx].cpu(), s=1, label="predictions", color="lightskyblue")
|
|
325
|
+
else:
|
|
326
|
+
y_sample = self.sample(x, sample_size=sample_size, expand_dim=False)
|
|
327
|
+
x_rep = x.repeat(sample_size, 1)
|
|
328
|
+
plt.scatter(x_rep[:,x_idx].cpu(), y_sample[:,y_idx].cpu(), s=1, label="samples", color="lightskyblue", alpha=alpha)
|
|
329
|
+
plt.legend(markerscale=2)
|
|
330
|
+
plt.ylim(ymin, ymax)
|
|
331
|
+
if x.shape[1] == 1:
|
|
332
|
+
plt.xlabel(r"$x$")
|
|
333
|
+
else:
|
|
334
|
+
plt.xlabel(r"$x_{}$".format(x_idx))
|
|
335
|
+
if y.shape[1] == 1:
|
|
336
|
+
plt.ylabel(r"$y$")
|
|
337
|
+
else:
|
|
338
|
+
plt.ylabel(r"$y_{}$".format(y_idx))
|
|
339
|
+
if save_dir is not None:
|
|
340
|
+
make_folder(save_dir)
|
|
341
|
+
plt.savefig(save_dir, bbox_inches="tight")
|
|
342
|
+
plt.close()
|
|
343
|
+
else:
|
|
344
|
+
plt.show()
|
|
345
|
+
|
|
346
|
+
|
|
347
|
+
def engression(x, y,
|
|
348
|
+
num_layer=2, hidden_dim=100, noise_dim=100,
|
|
349
|
+
lr=0.001, num_epoches=500, batch_size=None,
|
|
350
|
+
print_every_nepoch=100, print_times_per_epoch=1,
|
|
351
|
+
device="cpu", standardize=True,
|
|
352
|
+
verbose=True):
|
|
353
|
+
engressor = Engressor(in_dim=x.shape[1], out_dim=y.shape[1], num_layer=num_layer, hidden_dim=hidden_dim, noise_dim=noise_dim,
|
|
354
|
+
lr=lr, num_epoches=num_epoches, batch_size=batch_size, device=device, standardize=standardize)
|
|
355
|
+
engressor.train(x, y, num_epoches=num_epoches, batch_size=batch_size,
|
|
356
|
+
print_every_nepoch=print_every_nepoch, print_times_per_epoch=print_times_per_epoch,
|
|
357
|
+
standardize=standardize, verbose=verbose)
|
|
358
|
+
return engressor
|
|
@@ -0,0 +1,282 @@
|
|
|
1
|
+
# coming soon.
|
|
2
|
+
|
|
3
|
+
import torch
|
|
4
|
+
import numpy as np
|
|
5
|
+
import matplotlib.pyplot as plt
|
|
6
|
+
|
|
7
|
+
from .engression import Engressor
|
|
8
|
+
from .loss_func import *
|
|
9
|
+
from .models import StoNet
|
|
10
|
+
from .data.loader import make_dataloader
|
|
11
|
+
from .utils import *
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
class BaggedEngressor(object):
|
|
15
|
+
def __init__(self,
|
|
16
|
+
in_dim, out_dim, num_layer=2, hidden_dim=100, noise_dim=100,
|
|
17
|
+
lr=0.001, num_epoches=500, batch_size=None, device="cpu", standardize=True,
|
|
18
|
+
ensemble_size=10, val_loss_type="energy"):
|
|
19
|
+
"""Engressor with bagging.
|
|
20
|
+
|
|
21
|
+
Args:
|
|
22
|
+
in_dim (int): input dimension
|
|
23
|
+
out_dim (int): output dimension
|
|
24
|
+
num_layer (int, optional): number of layers. Defaults to 2.
|
|
25
|
+
hidden_dim (int, optional): number of neurons per layer. Defaults to 100.
|
|
26
|
+
noise_dim (int, optional): noise dimension. Defaults to 100.
|
|
27
|
+
lr (float, optional): learning rate. Defaults to 0.001.
|
|
28
|
+
num_epoches (int, optional): number of epoches. Defaults to 500.
|
|
29
|
+
batch_size (int, optional): batch size. Defaults to None, referring to the full batch.
|
|
30
|
+
device (str or torch.device, optional): device. Defaults to "cpu".
|
|
31
|
+
standardize (bool, optional): whether to standardize data. Defaults to True.
|
|
32
|
+
ensemble_size (int, optional): number of models for ensemble. Defaults to 10.
|
|
33
|
+
val_loss_type (str, optional): loss type for validation. Defaults to "energy". Choices: ["l1", "l2", "energy"].
|
|
34
|
+
"""
|
|
35
|
+
super().__init__()
|
|
36
|
+
self.lr = lr
|
|
37
|
+
self.num_epoches = num_epoches
|
|
38
|
+
self.batch_size = batch_size
|
|
39
|
+
self.ensemble_size = ensemble_size
|
|
40
|
+
self.num_models_for_each_sample = ensemble_size // 2
|
|
41
|
+
if isinstance(device, str):
|
|
42
|
+
device = torch.device(device)
|
|
43
|
+
self.device = device
|
|
44
|
+
self.standardize = standardize
|
|
45
|
+
self.x_mean = None
|
|
46
|
+
self.x_std = None
|
|
47
|
+
self.y_mean = None
|
|
48
|
+
self.y_std = None
|
|
49
|
+
self.split_mask = None
|
|
50
|
+
|
|
51
|
+
# Build ensemble models.
|
|
52
|
+
self.models = []
|
|
53
|
+
for _ in range(ensemble_size):
|
|
54
|
+
self.models.append(Engressor(in_dim, out_dim, num_layer, hidden_dim, noise_dim,
|
|
55
|
+
lr=lr, device=device, standardize=standardize))
|
|
56
|
+
self.val_loss_type = val_loss_type
|
|
57
|
+
|
|
58
|
+
def train_mode(self):
|
|
59
|
+
for i in range(self.ensemble_size):
|
|
60
|
+
self.models[i].train_mode()
|
|
61
|
+
|
|
62
|
+
def eval_mode(self):
|
|
63
|
+
for i in range(self.ensemble_size):
|
|
64
|
+
self.models[i].eval_mode()
|
|
65
|
+
|
|
66
|
+
def train(self, x, y, num_epoches=None, batch_size=None, standardize=True, val_loss_type="", val_sample_size=100, verbose=True):
|
|
67
|
+
"""_summary_
|
|
68
|
+
|
|
69
|
+
Args:
|
|
70
|
+
x (_type_): _description_
|
|
71
|
+
y (_type_): _description_
|
|
72
|
+
num_epoches (_type_, optional): _description_. Defaults to None.
|
|
73
|
+
batch_size (_type_, optional): _description_. Defaults to None.
|
|
74
|
+
standardize (bool, optional): _description_. Defaults to True.
|
|
75
|
+
- Standardization scheme for bagging:
|
|
76
|
+
for each model, all data are standardized using the mean and standard deviation of the training data for this model alone.
|
|
77
|
+
val_loss_type (str, optional): _description_. Defaults to "".
|
|
78
|
+
val_sample_size (int, optional): _description_. Defaults to 100.
|
|
79
|
+
verbose (bool, optional): _description_. Defaults to True.
|
|
80
|
+
"""
|
|
81
|
+
self.train_mode()
|
|
82
|
+
if num_epoches is None:
|
|
83
|
+
num_epoches = self.num_epoches
|
|
84
|
+
if batch_size is None:
|
|
85
|
+
batch_size = self.batch_size
|
|
86
|
+
if standardize:
|
|
87
|
+
self.standardize = standardize
|
|
88
|
+
|
|
89
|
+
x = vectorize(x)
|
|
90
|
+
y = vectorize(y)
|
|
91
|
+
x = x.to(self.device)
|
|
92
|
+
y = y.to(self.device)
|
|
93
|
+
|
|
94
|
+
# Mask matrix for splitting training and validation data.
|
|
95
|
+
data_size = x.shape[0]
|
|
96
|
+
rng = np.random.default_rng(21875667591346)
|
|
97
|
+
self.split_mask = torch.from_numpy(rng.multivariate_hypergeometric([1]*self.ensemble_size, self.num_models_for_each_sample, size=data_size))
|
|
98
|
+
|
|
99
|
+
# Training.
|
|
100
|
+
for i in range(self.ensemble_size):
|
|
101
|
+
model = self.models[i]
|
|
102
|
+
train_idx = self.split_mask[:,i]
|
|
103
|
+
x_train = x[train_idx]
|
|
104
|
+
y_train = y[train_idx]
|
|
105
|
+
x_val = x[~train_idx]
|
|
106
|
+
y_val = y[~train_idx]
|
|
107
|
+
model.train(x_train, y_train, num_epoches=num_epoches, batch_size=batch_size, standardize=self.standardize, verbose=False)
|
|
108
|
+
val_loss = model.eval_loss(x_val, y_val, loss_type="energy", sample_size=2)
|
|
109
|
+
train_loss = model.eval_loss(x_train, y_train, loss_type="energy", sample_size=2)
|
|
110
|
+
standardize_str = "(standardized)" if self.standardize else ""
|
|
111
|
+
print("[Model {}] train_loss{}: {:.4f}, val_loss{}: {:.4f}".format(i + 1, standardize_str, train_loss, standardize_str, val_loss))
|
|
112
|
+
|
|
113
|
+
# Validation.
|
|
114
|
+
val_loss_final = self.validate_bagged(x, y, loss_type=val_loss_type, sample_size=val_sample_size)
|
|
115
|
+
print("Final validation {} loss: {:.4f}".format(val_loss_type, val_loss_final))
|
|
116
|
+
|
|
117
|
+
def validate_bagged(self, x, y, loss_type="energy", sample_size=100):
|
|
118
|
+
"""Evaluate the bagged model on the validation data.
|
|
119
|
+
|
|
120
|
+
Args:
|
|
121
|
+
x (torch.Tensor): training data for predictors.
|
|
122
|
+
y (torch.Tensor): training data for responses.
|
|
123
|
+
loss_type (str, optional): type of the loss. Defaults to "energy".
|
|
124
|
+
sample_size (int, optional): _description_. Defaults to 100.
|
|
125
|
+
|
|
126
|
+
Returns:
|
|
127
|
+
float: loss of the bagged model on the validation data.
|
|
128
|
+
"""
|
|
129
|
+
self.eval_mode()
|
|
130
|
+
x = vectorize(x)
|
|
131
|
+
y = vectorize(y)
|
|
132
|
+
x = x.to(self.device)
|
|
133
|
+
y = y.to(self.device)
|
|
134
|
+
|
|
135
|
+
# Samples from the bagged model on the validation data.
|
|
136
|
+
sample_size_per_model = sample_size // self.num_models_for_each_sample
|
|
137
|
+
y_samples = torch.zeros((y.shape[0], y.shape[1], sample_size_per_model, self.ensemble_size), device=self.device)
|
|
138
|
+
for i in range(self.ensemble_size):
|
|
139
|
+
model = self.models[i]
|
|
140
|
+
val_idx = self.split_mask[:,i] == 0
|
|
141
|
+
y_samples[val_idx, :, :, i] = model.sample(x[val_idx], sample_size=sample_size_per_model)
|
|
142
|
+
val_idx = (1 - self.split_mask).nonzero(as_tuple=True)
|
|
143
|
+
y_samples = y_samples[val_idx[0], :, :, val_idx[1]]
|
|
144
|
+
y_samples = y_samples.reshape(y.shape[0], self.num_models_for_each_sample, y.shape[1], sample_size_per_model).permute(0, 2, 3, 1)
|
|
145
|
+
|
|
146
|
+
# Compute the loss (on the original scale).
|
|
147
|
+
if loss_type == "energy":
|
|
148
|
+
loss = energy_loss(y, y_samples, verbose=False)
|
|
149
|
+
elif loss_type == "l2":
|
|
150
|
+
y_pred = y_samples.mean(dim=len(y_samples.shape) - 1)
|
|
151
|
+
loss = (y - y_pred).pow(2).mean()
|
|
152
|
+
elif loss_type == "l1":
|
|
153
|
+
y_pred = y_samples.median(dim=len(y_samples.shape) - 1)
|
|
154
|
+
loss = (y - y_pred).abs().mean()
|
|
155
|
+
return loss.item()
|
|
156
|
+
|
|
157
|
+
def sample(self, x, sample_size=100):
|
|
158
|
+
self.eval_mode()
|
|
159
|
+
x = vectorize(x)
|
|
160
|
+
x = x.to(self.device)
|
|
161
|
+
sample_size_per_model = sample_size // self.ensemble_size
|
|
162
|
+
y_samples = []
|
|
163
|
+
for i in range(self.ensemble_size):
|
|
164
|
+
model = self.models[i]
|
|
165
|
+
if i == self.ensemble_size:
|
|
166
|
+
sample_size_per_model += sample_size % self.ensemble_size
|
|
167
|
+
y_samples.append(model.sample(x, sample_size=sample_size_per_model))
|
|
168
|
+
y_samples = torch.cat(y_samples, dim=-1)
|
|
169
|
+
return y_samples
|
|
170
|
+
|
|
171
|
+
def predict(self, x, target="mean", sample_size=100):
|
|
172
|
+
self.eval_mode()
|
|
173
|
+
x = vectorize(x)
|
|
174
|
+
x = x.to(self.device)
|
|
175
|
+
y_samples = self.sample(x, sample_size=sample_size)
|
|
176
|
+
if not isinstance(target, list):
|
|
177
|
+
target = [target]
|
|
178
|
+
y_pred = []
|
|
179
|
+
extremes = []
|
|
180
|
+
for t in target:
|
|
181
|
+
if t == "mean":
|
|
182
|
+
y_pred.append(y_samples.mean(dim=len(y_samples.shape) - 1))
|
|
183
|
+
else:
|
|
184
|
+
if t == "median":
|
|
185
|
+
t = 0.5
|
|
186
|
+
assert isinstance(t, float)
|
|
187
|
+
y_pred.append(y_samples.quantile(t, dim=len(y_samples.shape) - 1))
|
|
188
|
+
if min(t, 1 - t) * sample_size < 10:
|
|
189
|
+
extremes.append(t)
|
|
190
|
+
|
|
191
|
+
if len(extremes) > 0:
|
|
192
|
+
print("Warning: the estimate for quantiles at {} with a sample size of {} could be inaccurate. Please increase the `sample_size`.".format(extremes, sample_size))
|
|
193
|
+
|
|
194
|
+
if len(y_pred) == 1:
|
|
195
|
+
return y_pred[0]
|
|
196
|
+
else:
|
|
197
|
+
return y_pred
|
|
198
|
+
|
|
199
|
+
def eval_loss(self, x, y, loss_type="energy", sample_size=100):
|
|
200
|
+
"""Evaluate the bagged model.
|
|
201
|
+
|
|
202
|
+
Args:
|
|
203
|
+
x (_type_): _description_
|
|
204
|
+
y (_type_): _description_
|
|
205
|
+
loss_type (str, optional): _description_. Defaults to "energy".
|
|
206
|
+
sample_size (int, optional): _description_. Defaults to 100.
|
|
207
|
+
|
|
208
|
+
Returns:
|
|
209
|
+
_type_: _description_
|
|
210
|
+
"""
|
|
211
|
+
self.eval_mode()
|
|
212
|
+
x = vectorize(x)
|
|
213
|
+
y = vectorize(y)
|
|
214
|
+
x = x.to(self.device)
|
|
215
|
+
y = y.to(self.device)
|
|
216
|
+
if loss_type == "l2":
|
|
217
|
+
y_pred = self.predict(x, target="mean", sample_size=sample_size)
|
|
218
|
+
loss = (y - y_pred).pow(2).mean()
|
|
219
|
+
elif loss_type == "l1":
|
|
220
|
+
y_pred = self.predict(x, target=0.5, sample_size=sample_size)
|
|
221
|
+
loss = (y - y_pred).abs().mean()
|
|
222
|
+
else:
|
|
223
|
+
assert loss_type == "energy"
|
|
224
|
+
y_samples = self.sample(x, sample_size=sample_size, expand_dim=False)
|
|
225
|
+
loss = energy_loss(y, y_samples, verbose=False)
|
|
226
|
+
return loss.item()
|
|
227
|
+
|
|
228
|
+
def plot(self, x, y, x_idx=0, y_idx=0, target="mean", sample_size=100, save_dir=None,
|
|
229
|
+
alpha=0.8, ymin=None, ymax=None):
|
|
230
|
+
"""Plot true data and predictions.
|
|
231
|
+
|
|
232
|
+
Args:
|
|
233
|
+
x (torch.Tensor): data of predictors
|
|
234
|
+
y (torch.Tensor): data of responses
|
|
235
|
+
x_idx (int, optional): index of the predictor to plot (if there are multiple). Defaults to 0.
|
|
236
|
+
y_idx (int, optional): index of the response to plot (if there are multiple). Defaults to 0.
|
|
237
|
+
target (str or float, optional): target quantity. Defaults to "mean". Choice: ["mean", "median", "sample", float].
|
|
238
|
+
sample_size (int, optional): sample size used for estimation. Defaults to 100.
|
|
239
|
+
save_dir (str, optional): directory to save the plot. Defaults to None.
|
|
240
|
+
alpha (float, optional): transparency of the sampled data points. Defaults to 0.8.
|
|
241
|
+
ymin (float, optional): minimum value of y in the plot. Defaults to None.
|
|
242
|
+
ymax (float, optional): maximum value of y in the plot. Defaults to None.
|
|
243
|
+
"""
|
|
244
|
+
x = vectorize(x)
|
|
245
|
+
y = vectorize(y)
|
|
246
|
+
x = x.to(self.device)
|
|
247
|
+
y = y.to(self.device)
|
|
248
|
+
plt.scatter(x[:,x_idx].cpu(), y[:,y_idx].cpu(), s=1, label="true data", color="silver")
|
|
249
|
+
if target != "sample":
|
|
250
|
+
y_pred = self.predict(x, target=target, sample_size=sample_size)
|
|
251
|
+
plt.scatter(x[:,x_idx].cpu(), y_pred[:,y_idx].cpu(), s=1, label="predictions", color="lightskyblue")
|
|
252
|
+
else:
|
|
253
|
+
y_samples = self.sample(x, sample_size=sample_size, expand_dim=False)
|
|
254
|
+
x_rep = x.repeat(sample_size, 1)
|
|
255
|
+
plt.scatter(x_rep[:,x_idx].cpu(), y_samples[:,y_idx].cpu(), s=1, label="samples", color="lightskyblue", alpha=alpha)
|
|
256
|
+
plt.legend(markerscale=2)
|
|
257
|
+
plt.ylim(ymin, ymax)
|
|
258
|
+
if x.shape[1] == 1:
|
|
259
|
+
plt.xlabel(r"$x$")
|
|
260
|
+
else:
|
|
261
|
+
plt.xlabel(r"$x_{}$".format(x_idx))
|
|
262
|
+
if y.shape[1] == 1:
|
|
263
|
+
plt.ylabel(r"$y$")
|
|
264
|
+
else:
|
|
265
|
+
plt.ylabel(r"$y_{}$".format(y_idx))
|
|
266
|
+
if save_dir is not None:
|
|
267
|
+
make_folder(save_dir)
|
|
268
|
+
plt.savefig(save_dir, bbox_inches="tight")
|
|
269
|
+
plt.close()
|
|
270
|
+
else:
|
|
271
|
+
plt.show()
|
|
272
|
+
|
|
273
|
+
|
|
274
|
+
def engression_bagged(x, y,
|
|
275
|
+
num_layer=2, hidden_dim=100, noise_dim=100,
|
|
276
|
+
lr=0.001, num_epoches=500, batch_size=None, device="cpu",
|
|
277
|
+
nfolds=10, loss_type="energy"):
|
|
278
|
+
engressor = BaggedEngressor(in_dim=x.shape[1], out_dim=y.shape[1], num_layer=num_layer, hidden_dim=hidden_dim, noise_dim=noise_dim,
|
|
279
|
+
lr=lr, num_epoches=num_epoches, batch_size=batch_size, device=device, nfolds=nfolds)
|
|
280
|
+
engressor.train(x, y)
|
|
281
|
+
print("Average validation {} loss: {:.4f}".format(loss_type, np.mean(engressor.val_losses)))
|
|
282
|
+
return engressor
|
engression/loss_func.py
ADDED
|
@@ -0,0 +1,52 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
from .utils import vectorize
|
|
3
|
+
|
|
4
|
+
def energy_loss(x_true, x_est, beta=1, verbose=False):
|
|
5
|
+
"""Loss function based on the energy score.
|
|
6
|
+
|
|
7
|
+
Args:
|
|
8
|
+
x_true (torch.Tensor): iid samples from the true distribution.
|
|
9
|
+
x_est (list of torch.Tensor): a list of iid samples from the estimated distribution. ## todo: not a list!!!
|
|
10
|
+
beta (float): power parameter in the energy score.
|
|
11
|
+
verbose (bool): whether to return two terms of the loss.
|
|
12
|
+
|
|
13
|
+
Returns:
|
|
14
|
+
loss (torch.Tensor): energy loss.
|
|
15
|
+
"""
|
|
16
|
+
x_true = vectorize(x_true).unsqueeze(1)
|
|
17
|
+
if not isinstance(x_est, list):
|
|
18
|
+
x_est = list(torch.split(x_est, x_true.shape[0], dim=0))
|
|
19
|
+
m = len(x_est)
|
|
20
|
+
x_est = [vectorize(x_est[i]).unsqueeze(1) for i in range(m)]
|
|
21
|
+
x_est = torch.cat(x_est, dim=1)
|
|
22
|
+
|
|
23
|
+
s1 = torch.norm(x_est - x_true, 2, dim=2).pow(beta).mean()
|
|
24
|
+
s2 = torch.cdist(x_est, x_est, 2).pow(beta).mean() * m / (m - 1)
|
|
25
|
+
if verbose:
|
|
26
|
+
return torch.cat([(s1 - s2 / 2).reshape(1), s1.reshape(1), s2.reshape(1)], dim=0)
|
|
27
|
+
else:
|
|
28
|
+
return (s1 - s2 / 2)
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def energy_loss_two_sample(x0, x, xp, beta=1, verbose=False):
|
|
32
|
+
"""Loss function based on the energy score (estimated based on two samples).
|
|
33
|
+
|
|
34
|
+
Args:
|
|
35
|
+
x0 (torch.Tensor): iid samples from the true distribution.
|
|
36
|
+
x (torch.Tensor): iid samples from the estimated distribution.
|
|
37
|
+
xp (torch.Tensor): iid samples from the estimated distribution.
|
|
38
|
+
beta (float): power parameter in the energy score.
|
|
39
|
+
verbose (bool): whether to return two terms of the loss.
|
|
40
|
+
|
|
41
|
+
Returns:
|
|
42
|
+
loss (torch.Tensor): energy loss.
|
|
43
|
+
"""
|
|
44
|
+
x0 = vectorize(x0)
|
|
45
|
+
x = vectorize(x)
|
|
46
|
+
xp = vectorize(xp)
|
|
47
|
+
s1 = torch.norm(x - x0, 2, dim=1).pow(beta).mean() / 2 + torch.norm(xp - x0, 2, dim=1).pow(beta).mean() / 2
|
|
48
|
+
s2 = torch.norm(x - xp, 2, dim=1).pow(beta).mean()
|
|
49
|
+
if verbose:
|
|
50
|
+
return torch.cat([(s1 - s2/2).reshape(1), s1.reshape(1), s2.reshape(1)], dim=0)
|
|
51
|
+
else:
|
|
52
|
+
return (s1 - s2/2)
|
engression/models.py
ADDED
|
@@ -0,0 +1,179 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
import torch.nn as nn
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
class StoLayer(nn.Module):
|
|
6
|
+
"""A stochastic layer.
|
|
7
|
+
|
|
8
|
+
Args:
|
|
9
|
+
in_dim (int): input dimension
|
|
10
|
+
out_dim (int): output dimension
|
|
11
|
+
noise_dim (int, optional): noise dimension. Defaults to 100.
|
|
12
|
+
"""
|
|
13
|
+
def __init__(self, in_dim, out_dim, noise_dim=100):
|
|
14
|
+
super().__init__()
|
|
15
|
+
self.in_dim = in_dim
|
|
16
|
+
self.out_dim = out_dim
|
|
17
|
+
self.noise_dim = noise_dim
|
|
18
|
+
layer = [
|
|
19
|
+
nn.Linear(in_dim + noise_dim, out_dim),
|
|
20
|
+
nn.BatchNorm1d(out_dim),
|
|
21
|
+
nn.ReLU(inplace=True),
|
|
22
|
+
]
|
|
23
|
+
self.layer = nn.Sequential(*layer)
|
|
24
|
+
|
|
25
|
+
def forward(self, x, inject_noise=True):
|
|
26
|
+
eps = torch.randn(x.size(0), self.noise_dim, device=x.device) if inject_noise else torch.zeros(x.size(0), self.noise_dim, device=x.device)
|
|
27
|
+
x = torch.cat([x, eps], dim=1)
|
|
28
|
+
return self.layer(x)
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
class StoNet(nn.Module):
|
|
32
|
+
"""Stochastic neural network.
|
|
33
|
+
|
|
34
|
+
Args:
|
|
35
|
+
in_dim (int): input dimension
|
|
36
|
+
out_dim (int): output dimension
|
|
37
|
+
num_layer (int, optional): number of layers. Defaults to 2.
|
|
38
|
+
hidden_dim (int, optional): number of neurons per layer. Defaults to 100.
|
|
39
|
+
noise_dim (int, optional): noise dimension. Defaults to 100.
|
|
40
|
+
"""
|
|
41
|
+
def __init__(self, in_dim, out_dim, num_layer=2, hidden_dim=100, noise_dim=100):
|
|
42
|
+
super().__init__()
|
|
43
|
+
self.in_dim = in_dim
|
|
44
|
+
self.out_dim = out_dim
|
|
45
|
+
self.num_layer = num_layer
|
|
46
|
+
self.hidden_dim = hidden_dim
|
|
47
|
+
self.noise_dim = noise_dim
|
|
48
|
+
|
|
49
|
+
self.input_layer = StoLayer(in_dim, hidden_dim, noise_dim)
|
|
50
|
+
if num_layer > 2:
|
|
51
|
+
inter_layer = [StoLayer(hidden_dim, hidden_dim, noise_dim)]
|
|
52
|
+
for i in range(num_layer - 3):
|
|
53
|
+
inter_layer.append(StoLayer(hidden_dim, hidden_dim, noise_dim))
|
|
54
|
+
self.inter_layer = nn.Sequential(*inter_layer)
|
|
55
|
+
self.out_layer = nn.Linear(hidden_dim, out_dim)
|
|
56
|
+
|
|
57
|
+
def predict(self, x, target=["mean"], sample_size=100):
|
|
58
|
+
"""Point prediction.
|
|
59
|
+
|
|
60
|
+
Args:
|
|
61
|
+
x (torch.Tensor): _description_
|
|
62
|
+
target (str or float or list, optional): single-valued functional to predict. float refers to the quantiles. Defaults to ["mean"].
|
|
63
|
+
sample_size (int, optional): sample sizes for each x. Defaults to 100.
|
|
64
|
+
|
|
65
|
+
Returns:
|
|
66
|
+
torch.Tensor or list of torch.Tensor: point predictions
|
|
67
|
+
- [:,:,i] gives the i-th sample of all x.
|
|
68
|
+
- [i,:,:] gives all samples of x_i.
|
|
69
|
+
|
|
70
|
+
Here we do not call `sample` but directly call `forward`.
|
|
71
|
+
"""
|
|
72
|
+
# eval_size = x.size(0)
|
|
73
|
+
# # Sampling
|
|
74
|
+
# with torch.no_grad():
|
|
75
|
+
# x_rep = torch.repeat_interleave(x, (sample_size * torch.ones(eval_size, device=x.device)).long(), dim=0)
|
|
76
|
+
# y_pred = self.forward(x=x_rep).detach()
|
|
77
|
+
# y_pred = list(torch.split(y_pred, sample_size))
|
|
78
|
+
|
|
79
|
+
# if not isinstance(target, list):
|
|
80
|
+
# target = [target]
|
|
81
|
+
# results = []
|
|
82
|
+
# extremes = []
|
|
83
|
+
# for t in target:
|
|
84
|
+
# if t == "mean":
|
|
85
|
+
# results.append(torch.cat([y_pred[i].mean().unsqueeze(0) for i in range(eval_size)], dim=0).unsqueeze(1))
|
|
86
|
+
# else:
|
|
87
|
+
# if t == "median":
|
|
88
|
+
# t = 0.5
|
|
89
|
+
# assert isinstance(t, float)
|
|
90
|
+
# results.append(torch.cat([y_pred[i].quantile(t).unsqueeze(0) for i in range(eval_size)], dim=0).unsqueeze(1))
|
|
91
|
+
# if min(t, 1-t) * sample_size < 10:
|
|
92
|
+
# extremes.append(t)
|
|
93
|
+
|
|
94
|
+
samples = self.sample(x=x, sample_size=sample_size, expand_dim=True)
|
|
95
|
+
if not isinstance(target, list):
|
|
96
|
+
target = [target]
|
|
97
|
+
results = []
|
|
98
|
+
extremes = []
|
|
99
|
+
for t in target:
|
|
100
|
+
if t == "mean":
|
|
101
|
+
results.append(samples.mean(dim=len(samples.shape) - 1))
|
|
102
|
+
else:
|
|
103
|
+
if t == "median":
|
|
104
|
+
t = 0.5
|
|
105
|
+
assert isinstance(t, float)
|
|
106
|
+
results.append(samples.quantile(t, dim=len(samples.shape) - 1))
|
|
107
|
+
if min(t, 1 - t) * sample_size < 10:
|
|
108
|
+
extremes.append(t)
|
|
109
|
+
|
|
110
|
+
if len(extremes) > 0:
|
|
111
|
+
print("Warning: the estimate for quantiles at {} with a sample size of {} could be inaccurate. Please increase the `sample_size`.".format(extremes, sample_size))
|
|
112
|
+
|
|
113
|
+
if len(results) == 1:
|
|
114
|
+
return results[0]
|
|
115
|
+
else:
|
|
116
|
+
return results
|
|
117
|
+
|
|
118
|
+
def sample(self, x, sample_size=100, expand_dim=True):
|
|
119
|
+
"""Sample new response data.
|
|
120
|
+
|
|
121
|
+
Args:
|
|
122
|
+
x (torch.Tensor): new data of predictors of shape [data_size, covariate_dim]
|
|
123
|
+
sample_size (int, optional): new sample size. Defaults to 100.
|
|
124
|
+
expand_dim (bool, optional): whether to expand the sample dimension. Defaults to True.
|
|
125
|
+
|
|
126
|
+
Returns:
|
|
127
|
+
torch.Tensor of shape (data_size, response_dim, sample_size), where response_dim could have multiple channels.
|
|
128
|
+
"""
|
|
129
|
+
eval_size = x.size(0)
|
|
130
|
+
with torch.no_grad():
|
|
131
|
+
# x_rep = torch.repeat_interleave(x, (sample_size * torch.ones(eval_size, device=x.device)).long(), dim=0)
|
|
132
|
+
x_rep = x.repeat(sample_size, 1)
|
|
133
|
+
samples = self.forward(x=x_rep).detach()
|
|
134
|
+
if not expand_dim:
|
|
135
|
+
return samples
|
|
136
|
+
else:
|
|
137
|
+
expand_dim = len(samples.shape)
|
|
138
|
+
samples = samples.unsqueeze(expand_dim)
|
|
139
|
+
samples = list(torch.split(samples, eval_size))
|
|
140
|
+
samples = torch.cat(samples, dim=expand_dim)
|
|
141
|
+
return samples
|
|
142
|
+
# without expanding dimensions:
|
|
143
|
+
# samples.reshape(-1, *samples.shape[1:-1])
|
|
144
|
+
|
|
145
|
+
def forward(self, x, inject_noise=True):
|
|
146
|
+
x = self.input_layer(x, inject_noise)
|
|
147
|
+
for i in range(self.num_layer - 2):
|
|
148
|
+
x = self.inter_layer[i](x, inject_noise)
|
|
149
|
+
x = self.out_layer(x)
|
|
150
|
+
return x
|
|
151
|
+
|
|
152
|
+
|
|
153
|
+
class Net(nn.Module):
|
|
154
|
+
"""Deterministic neural network.
|
|
155
|
+
|
|
156
|
+
Args:
|
|
157
|
+
in_dim (int, optional): input dimension. Defaults to 1.
|
|
158
|
+
out_dim (int, optional): output dimension. Defaults to 1.
|
|
159
|
+
num_layer (int, optional): number of layers. Defaults to 2.
|
|
160
|
+
hidden_dim (int, optional): number of neurons per layer. Defaults to 100.
|
|
161
|
+
"""
|
|
162
|
+
def __init__(self, in_dim=1, out_dim=1, num_layer=2, hidden_dim=100):
|
|
163
|
+
super().__init__()
|
|
164
|
+
net = [
|
|
165
|
+
nn.Linear(in_dim, hidden_dim),
|
|
166
|
+
nn.BatchNorm1d(hidden_dim),
|
|
167
|
+
nn.ReLU(inplace=True),
|
|
168
|
+
]
|
|
169
|
+
for _ in range(num_layer - 2):
|
|
170
|
+
net += [
|
|
171
|
+
nn.Linear(hidden_dim, hidden_dim),
|
|
172
|
+
nn.BatchNorm1d(hidden_dim),
|
|
173
|
+
nn.ReLU(inplace=True),
|
|
174
|
+
]
|
|
175
|
+
net.append(nn.Linear(hidden_dim, out_dim))
|
|
176
|
+
self.net = nn.Sequential(*net)
|
|
177
|
+
|
|
178
|
+
def forward(self, x):
|
|
179
|
+
return self.net(x)
|
engression/utils.py
ADDED
|
@@ -0,0 +1,65 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import torch
|
|
3
|
+
|
|
4
|
+
def vectorize(x, multichannel=False):
|
|
5
|
+
"""Vectorize data in any shape.
|
|
6
|
+
|
|
7
|
+
Args:
|
|
8
|
+
x (torch.Tensor): input data
|
|
9
|
+
multichannel (bool, optional): whether to keep the multiple channels (in the second dimension). Defaults to False.
|
|
10
|
+
|
|
11
|
+
Returns:
|
|
12
|
+
torch.Tensor: data of shape (sample_size, dimension) or (sample_size, num_channel, dimension) if multichannel is True.
|
|
13
|
+
"""
|
|
14
|
+
if len(x.shape) == 1:
|
|
15
|
+
return x.unsqueeze(1)
|
|
16
|
+
if len(x.shape) == 2:
|
|
17
|
+
return x
|
|
18
|
+
else:
|
|
19
|
+
if not multichannel: # one channel
|
|
20
|
+
return x.reshape(x.shape[0], -1)
|
|
21
|
+
else: # multi-channel
|
|
22
|
+
return x.reshape(x.shape[0], x.shape[1], -1)
|
|
23
|
+
|
|
24
|
+
def cor(x, y):
|
|
25
|
+
"""Compute the correlation between two signals.
|
|
26
|
+
|
|
27
|
+
Args:
|
|
28
|
+
x (torch.Tensor): input data
|
|
29
|
+
y (torch.Tensor): input data
|
|
30
|
+
|
|
31
|
+
Returns:
|
|
32
|
+
torch.Tensor: correlation between x and y
|
|
33
|
+
"""
|
|
34
|
+
x = vectorize(x)
|
|
35
|
+
y = vectorize(y)
|
|
36
|
+
x = x - x.mean(0)
|
|
37
|
+
y = y - y.mean(0)
|
|
38
|
+
return ((x * y).mean()) / (x.std(unbiased=False) * y.std(unbiased=False))
|
|
39
|
+
|
|
40
|
+
def make_folder(name):
|
|
41
|
+
"""Make a folder.
|
|
42
|
+
|
|
43
|
+
Args:
|
|
44
|
+
name (str): folder name.
|
|
45
|
+
"""
|
|
46
|
+
if not os.path.exists(name):
|
|
47
|
+
print('Creating folder: {}'.format(name))
|
|
48
|
+
os.makedirs(name)
|
|
49
|
+
|
|
50
|
+
def check_for_gpu(device):
|
|
51
|
+
"""Check if a CUDA device is available.
|
|
52
|
+
|
|
53
|
+
Args:
|
|
54
|
+
device (torch.device): current set device.
|
|
55
|
+
"""
|
|
56
|
+
if device == torch.device('cuda'):
|
|
57
|
+
if torch.cuda.is_available():
|
|
58
|
+
print("GPU is available, running on GPU.\n")
|
|
59
|
+
else:
|
|
60
|
+
print("GPU is NOT available, running instead on CPU.\n")
|
|
61
|
+
else:
|
|
62
|
+
if torch.cuda.is_available():
|
|
63
|
+
print("Warning: You have a CUDA device, so you may consider using GPU for potential acceleration\n by setting device to 'cuda'.\n")
|
|
64
|
+
else:
|
|
65
|
+
print("Running on CPU.\n")
|
|
@@ -0,0 +1,47 @@
|
|
|
1
|
+
Metadata-Version: 2.1
|
|
2
|
+
Name: engression
|
|
3
|
+
Version: 0.0.0.dev0
|
|
4
|
+
Summary: Engression
|
|
5
|
+
Home-page: https://github.com/xwshen51/engression
|
|
6
|
+
Author: Xinwei Shen
|
|
7
|
+
Author-email: xinwei.shen@stat.math.ethz.ch
|
|
8
|
+
Classifier: Programming Language :: Python :: 3
|
|
9
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
10
|
+
Classifier: Operating System :: OS Independent
|
|
11
|
+
Description-Content-Type: text/markdown
|
|
12
|
+
Requires-Dist: torch
|
|
13
|
+
Requires-Dist: numpy
|
|
14
|
+
Requires-Dist: matplotlib
|
|
15
|
+
|
|
16
|
+
# Engression
|
|
17
|
+
|
|
18
|
+
Engression is a nonlinear regression methodology proposed in the paper "[*Engression: Extrapolation for Nonlinear Regression?*]()" by Xinwei Shen and Nicolai Meinshausen.
|
|
19
|
+
This directory contains the Python implementations of engression.
|
|
20
|
+
|
|
21
|
+
## Installation
|
|
22
|
+
The latest release of the Python package can be installed through pip:
|
|
23
|
+
```sh
|
|
24
|
+
pip install engression
|
|
25
|
+
```
|
|
26
|
+
|
|
27
|
+
The development version can be installed from github
|
|
28
|
+
|
|
29
|
+
```sh
|
|
30
|
+
$ git clone https://github.com/xwshen51/engression.git # Download the package
|
|
31
|
+
$ cd engression
|
|
32
|
+
$ pip install -r requirements.txt # Install the requirements
|
|
33
|
+
$ python setup.py install develop --user
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
## Usage Example
|
|
38
|
+
|
|
39
|
+
### Python
|
|
40
|
+
```python
|
|
41
|
+
import engression
|
|
42
|
+
|
|
43
|
+
```
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
## Contact information
|
|
47
|
+
If you meet any problems with the code, please submit an issue or contact Xinwei Shen (`xinwei.shen@stat.math.ethz.ch`).
|
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
engression/__init__.py,sha256=aS5GIvTRA2rpCrZOExqu6E9SvGcBZnI3E_VCsEqpFbc,365
|
|
2
|
+
engression/engression.py,sha256=PAhm9uE95IDRI-mr_42iL3Ztd5poHBfmmHgWWswktCM,16497
|
|
3
|
+
engression/engression_bagged.py,sha256=-_czVMqRWD2pJhHv0sZYfC4L95M9EwYwRskMiDbuXZ0,12725
|
|
4
|
+
engression/loss_func.py,sha256=aE5KpdhARtChHUnJ-7knQrpG6vBFZdiMjg2Rd_c5Qcg,2026
|
|
5
|
+
engression/models.py,sha256=NFNkhDNCO5Sl2C0RkKY2K8Eosh-_L75GRMx5ZGYY6tM,6999
|
|
6
|
+
engression/utils.py,sha256=2fwA-DT4wjPzGZwGOaQfB4yYVyGhLfJJTAPGUZYysxU,1910
|
|
7
|
+
engression/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
8
|
+
engression/data/loader.py,sha256=cITbaS5VH4hCWihCKnJWpX-xVS7glhwhDtmj1atWHUk,1980
|
|
9
|
+
engression/data/simulator.py,sha256=i5wONmiQUxymg_lsVuJGB3q0-ZjFp6jzqFsXc-mIl3Y,2673
|
|
10
|
+
engression-0.0.0.dev0.dist-info/METADATA,sha256=XFOKrupfz-3ozN3-p0WSNq7tA61zO1Qk48v3ISzF6OY,1295
|
|
11
|
+
engression-0.0.0.dev0.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
|
|
12
|
+
engression-0.0.0.dev0.dist-info/top_level.txt,sha256=N0yRTczPTx-rOI-BgXZrjo8FB422Z2eH7uB0U99tZMw,11
|
|
13
|
+
engression-0.0.0.dev0.dist-info/RECORD,,
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
engression
|