Sped up sampling 20x
This commit is contained in:
Victor Mylle
@@ -13,6 +13,49 @@ from tqdm import tqdm
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
|
||||
def sample_from_dist(quantiles, output_values):
|
||||
# both to numpy
|
||||
quantiles = quantiles.cpu().numpy()
|
||||
|
||||
if isinstance(output_values, torch.Tensor):
|
||||
output_values = output_values.cpu().numpy()
|
||||
|
||||
reshaped_values = output_values.reshape(-1, len(quantiles))
|
||||
|
||||
uniform_random_numbers = np.random.uniform(0, 1, (reshaped_values.shape[0], 1000))
|
||||
|
||||
idx_below = np.searchsorted(quantiles, uniform_random_numbers, side="right") - 1
|
||||
idx_above = np.clip(idx_below + 1, 0, len(quantiles) - 1)
|
||||
|
||||
# handle edge case where idx_below is -1
|
||||
idx_below = np.clip(idx_below, 0, len(quantiles) - 1)
|
||||
|
||||
y_below = reshaped_values[np.arange(reshaped_values.shape[0])[:, None], idx_below]
|
||||
y_above = reshaped_values[np.arange(reshaped_values.shape[0])[:, None], idx_above]
|
||||
|
||||
# Calculate the slopes for interpolation
|
||||
x_below = quantiles[idx_below]
|
||||
x_above = quantiles[idx_above]
|
||||
|
||||
# Interpolate
|
||||
# Ensure all variables are NumPy arrays
|
||||
x_below_np = x_below.cpu().numpy() if isinstance(x_below, torch.Tensor) else x_below
|
||||
x_above_np = x_above.cpu().numpy() if isinstance(x_above, torch.Tensor) else x_above
|
||||
y_below_np = y_below.cpu().numpy() if isinstance(y_below, torch.Tensor) else y_below
|
||||
y_above_np = y_above.cpu().numpy() if isinstance(y_above, torch.Tensor) else y_above
|
||||
|
||||
# Compute slopes for interpolation
|
||||
slopes_np = (y_above_np - y_below_np) / (
|
||||
np.clip(x_above_np - x_below_np, 1e-6, np.inf)
|
||||
)
|
||||
|
||||
# Perform the interpolation
|
||||
new_samples = y_below_np + slopes_np * (uniform_random_numbers - x_below_np)
|
||||
|
||||
# Return the mean of the samples
|
||||
return np.mean(new_samples, axis=1)
|
||||
|
||||
|
||||
class AutoRegressiveQuantileTrainer(AutoRegressiveTrainer):
|
||||
def __init__(
|
||||
self,
|
||||
@@ -46,19 +89,26 @@ class AutoRegressiveQuantileTrainer(AutoRegressiveTrainer):
|
||||
}
|
||||
|
||||
with torch.no_grad():
|
||||
total_amount_samples = len(dataloader.dataset) - 95
|
||||
total_samples = len(dataloader.dataset) - 96
|
||||
batches = 0
|
||||
for _, _, idx_batch in dataloader:
|
||||
idx_batch = [idx for idx in idx_batch if idx < total_samples]
|
||||
|
||||
for idx in tqdm(range(total_amount_samples)):
|
||||
_, outputs, samples, targets = self.auto_regressive(dataloader, idx)
|
||||
if len(idx_batch) == 0:
|
||||
continue
|
||||
|
||||
_, outputs, samples, targets = self.auto_regressive(
|
||||
dataloader.dataset, idx_batch=idx_batch
|
||||
)
|
||||
|
||||
samples = samples.to(self.device)
|
||||
outputs = outputs.to(self.device)
|
||||
targets = targets.to(self.device)
|
||||
|
||||
inversed_samples = self.data_processor.inverse_transform(samples)
|
||||
inversed_targets = self.data_processor.inverse_transform(targets)
|
||||
inversed_outputs = self.data_processor.inverse_transform(outputs)
|
||||
|
||||
outputs = outputs.to(self.device)
|
||||
targets = targets.to(self.device)
|
||||
samples = samples.to(self.device)
|
||||
|
||||
inversed_samples = inversed_samples.to(self.device)
|
||||
inversed_targets = inversed_targets.to(self.device)
|
||||
inversed_outputs = inversed_outputs.to(self.device)
|
||||
@@ -66,10 +116,10 @@ class AutoRegressiveQuantileTrainer(AutoRegressiveTrainer):
|
||||
for metric in self.metrics_to_track:
|
||||
if metric.__class__ != PinballLoss and metric.__class__ != CRPSLoss:
|
||||
transformed_metrics[metric.__class__.__name__] += metric(
|
||||
samples, targets
|
||||
samples, targets.squeeze(-1)
|
||||
)
|
||||
metrics[metric.__class__.__name__] += metric(
|
||||
inversed_samples, inversed_targets
|
||||
inversed_samples, inversed_targets.squeeze(-1)
|
||||
)
|
||||
else:
|
||||
transformed_metrics[metric.__class__.__name__] += metric(
|
||||
@@ -78,10 +128,11 @@ class AutoRegressiveQuantileTrainer(AutoRegressiveTrainer):
|
||||
metrics[metric.__class__.__name__] += metric(
|
||||
inversed_outputs, inversed_targets
|
||||
)
|
||||
batches += 1
|
||||
|
||||
for metric in self.metrics_to_track:
|
||||
metrics[metric.__class__.__name__] /= total_amount_samples
|
||||
transformed_metrics[metric.__class__.__name__] /= total_amount_samples
|
||||
metrics[metric.__class__.__name__] /= batches
|
||||
transformed_metrics[metric.__class__.__name__] /= batches
|
||||
|
||||
for metric_name, metric_value in metrics.items():
|
||||
if PinballLoss.__name__ in metric_name:
|
||||
@@ -97,7 +148,14 @@ class AutoRegressiveQuantileTrainer(AutoRegressiveTrainer):
|
||||
)
|
||||
task.get_logger().report_single_value(name=name, value=metric_value)
|
||||
|
||||
def get_plot(self, current_day, next_day, predictions, show_legend: bool = True):
|
||||
def get_plot(
|
||||
self,
|
||||
current_day,
|
||||
next_day,
|
||||
predictions,
|
||||
show_legend: bool = True,
|
||||
retransform: bool = True,
|
||||
):
|
||||
fig = go.Figure()
|
||||
|
||||
# Convert to numpy for plotting
|
||||
@@ -105,6 +163,11 @@ class AutoRegressiveQuantileTrainer(AutoRegressiveTrainer):
|
||||
next_day_np = next_day.view(-1).cpu().numpy()
|
||||
predictions_np = predictions.cpu().numpy()
|
||||
|
||||
if retransform:
|
||||
current_day_np = self.data_processor.inverse_transform(current_day_np)
|
||||
next_day_np = self.data_processor.inverse_transform(next_day_np)
|
||||
predictions_np = self.data_processor.inverse_transform(predictions_np)
|
||||
|
||||
# Add traces for current and next day
|
||||
fig.add_trace(go.Scatter(x=np.arange(96), y=current_day_np, name="Current Day"))
|
||||
fig.add_trace(go.Scatter(x=96 + np.arange(96), y=next_day_np, name="Next Day"))
|
||||
@@ -127,86 +190,68 @@ class AutoRegressiveQuantileTrainer(AutoRegressiveTrainer):
|
||||
|
||||
return fig
|
||||
|
||||
def auto_regressive(self, data_loader, idx, sequence_length: int = 96):
|
||||
self.model.eval()
|
||||
target_full = []
|
||||
predictions_sampled = []
|
||||
predictions_full = []
|
||||
|
||||
prev_features, target = data_loader.dataset[idx]
|
||||
def auto_regressive(self, dataset, idx_batch, sequence_length: int = 96):
|
||||
prev_features, targets = dataset.get_batch(idx_batch)
|
||||
prev_features = prev_features.to(self.device)
|
||||
targets = targets.to(self.device)
|
||||
|
||||
initial_sequence = prev_features[:96]
|
||||
initial_sequence = prev_features[:, :96]
|
||||
|
||||
target_full.append(target)
|
||||
target_full = targets[:, 0].unsqueeze(1) # (batch_size, 1)
|
||||
with torch.no_grad():
|
||||
prediction = self.model(prev_features.unsqueeze(0))
|
||||
predictions_full.append(prediction.squeeze(0))
|
||||
|
||||
# sample from the distribution
|
||||
sample = self.sample_from_dist(
|
||||
self.quantiles.cpu(), prediction.squeeze(-1).cpu().numpy()
|
||||
)
|
||||
predictions_sampled.append(sample)
|
||||
new_predictions_full = self.model(prev_features) # (batch_size, quantiles)
|
||||
samples = (
|
||||
torch.tensor(sample_from_dist(self.quantiles, new_predictions_full))
|
||||
.unsqueeze(1)
|
||||
.to(self.device)
|
||||
) # (batch_size, 1)
|
||||
predictions_samples = samples
|
||||
predictions_full = new_predictions_full.unsqueeze(1)
|
||||
|
||||
for i in range(sequence_length - 1):
|
||||
new_features = torch.cat(
|
||||
(prev_features[1:96].cpu(), torch.tensor([predictions_sampled[-1]])),
|
||||
dim=0,
|
||||
)
|
||||
(prev_features[:, 1:96], samples), dim=1
|
||||
) # (batch_size, 96)
|
||||
|
||||
new_features = new_features.float()
|
||||
|
||||
# get the other needed features
|
||||
other_features, new_target = data_loader.dataset.random_day_autoregressive(
|
||||
idx + i + 1
|
||||
)
|
||||
other_features, new_targets = dataset.get_batch_autoregressive(
|
||||
np.array(idx_batch) + i + 1
|
||||
) # (batch_size, new_features)
|
||||
|
||||
if other_features is not None:
|
||||
prev_features = torch.cat((new_features, other_features), dim=0)
|
||||
prev_features = torch.cat(
|
||||
new_features, other_features, dim=1
|
||||
) # (batch_size, 96 + new_features)
|
||||
else:
|
||||
prev_features = new_features
|
||||
|
||||
# add target to target_full
|
||||
target_full.append(new_target)
|
||||
target_full = torch.cat(
|
||||
(target_full, new_targets.to(self.device)), dim=1
|
||||
) # (batch_size, sequence_length)
|
||||
|
||||
# predict
|
||||
with torch.no_grad():
|
||||
prediction = self.model(prev_features.unsqueeze(0).to(self.device))
|
||||
predictions_full.append(prediction.squeeze(0))
|
||||
new_predictions_full = self.model(
|
||||
prev_features
|
||||
) # (batch_size, quantiles)
|
||||
predictions_full = torch.cat(
|
||||
(predictions_full, new_predictions_full.unsqueeze(1)), dim=1
|
||||
) # (batch_size, sequence_length, quantiles)
|
||||
|
||||
# sample from the distribution
|
||||
sample = self.sample_from_dist(
|
||||
self.quantiles.cpu(), prediction.squeeze(-1).cpu().numpy()
|
||||
)
|
||||
predictions_sampled.append(sample)
|
||||
samples = (
|
||||
torch.tensor(sample_from_dist(self.quantiles, new_predictions_full))
|
||||
.unsqueeze(-1)
|
||||
.to(self.device)
|
||||
) # (batch_size, 1)
|
||||
predictions_samples = torch.cat((predictions_samples, samples), dim=1)
|
||||
|
||||
return (
|
||||
initial_sequence.cpu(),
|
||||
torch.stack(predictions_full).cpu(),
|
||||
torch.tensor(predictions_sampled).reshape(-1, 1),
|
||||
torch.stack(target_full).cpu(),
|
||||
initial_sequence,
|
||||
predictions_full,
|
||||
predictions_samples,
|
||||
target_full.unsqueeze(-1),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def sample_from_dist(quantiles, output_values):
|
||||
# Interpolate the inverse CDF
|
||||
inverse_cdf = interp1d(
|
||||
quantiles,
|
||||
output_values,
|
||||
kind="linear",
|
||||
bounds_error=False,
|
||||
fill_value="extrapolate",
|
||||
)
|
||||
|
||||
# generate one random uniform number
|
||||
uniform_random_numbers = np.random.uniform(0, 1, 1000)
|
||||
|
||||
# Apply the inverse CDF to the uniform random numbers
|
||||
samples = inverse_cdf(uniform_random_numbers)
|
||||
|
||||
# Return the mean of the samples
|
||||
return np.mean(samples)
|
||||
|
||||
def plot_quantile_percentages(
|
||||
self, task, data_loader, train: bool = True, iteration: int = None
|
||||
):
|
||||
@@ -214,7 +259,7 @@ class AutoRegressiveQuantileTrainer(AutoRegressiveTrainer):
|
||||
quantile_counter = {q: 0 for q in self.quantiles.cpu().numpy()}
|
||||
|
||||
with torch.no_grad():
|
||||
for inputs, targets in data_loader:
|
||||
for inputs, targets, _ in data_loader:
|
||||
inputs = inputs.to("cuda")
|
||||
output = self.model(inputs)
|
||||
|
||||
@@ -302,23 +347,6 @@ class NonAutoRegressiveQuantileRegression(Trainer):
|
||||
debug=debug,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def sample_from_dist(quantiles, output_values):
|
||||
reshaped_values = output_values.reshape(-1, len(quantiles))
|
||||
samples = []
|
||||
for row in reshaped_values:
|
||||
inverse_cdf = interp1d(
|
||||
quantiles,
|
||||
row,
|
||||
kind="linear",
|
||||
bounds_error=False,
|
||||
fill_value="extrapolate",
|
||||
)
|
||||
uniform_random_numbers = np.random.uniform(0, 1, 1000)
|
||||
new_samples = inverse_cdf(uniform_random_numbers)
|
||||
samples.append(np.mean(new_samples))
|
||||
return np.array(samples)
|
||||
|
||||
def log_final_metrics(self, task, dataloader, train: bool = True):
|
||||
metrics = {metric.__class__.__name__: 0.0 for metric in self.metrics_to_track}
|
||||
transformed_metrics = {
|
||||
@@ -326,12 +354,12 @@ class NonAutoRegressiveQuantileRegression(Trainer):
|
||||
}
|
||||
|
||||
with torch.no_grad():
|
||||
for inputs, targets in dataloader:
|
||||
for inputs, targets, _ in dataloader:
|
||||
inputs, targets = inputs.to(self.device), targets.to(self.device)
|
||||
|
||||
outputs = self.model(inputs)
|
||||
outputted_samples = [
|
||||
self.sample_from_dist(self.quantiles.cpu(), output.cpu().numpy())
|
||||
sample_from_dist(self.quantiles.cpu(), output.cpu().numpy())
|
||||
for output in outputs
|
||||
]
|
||||
|
||||
@@ -359,10 +387,10 @@ class NonAutoRegressiveQuantileRegression(Trainer):
|
||||
)
|
||||
else:
|
||||
transformed_metrics[metric.__class__.__name__] += metric(
|
||||
outputs, targets
|
||||
outputs, targets.unsqueeze(-1)
|
||||
)
|
||||
metrics[metric.__class__.__name__] += metric(
|
||||
inversed_outputs, inversed_targets
|
||||
inversed_outputs, inversed_targets.unsqueeze(-1)
|
||||
)
|
||||
|
||||
for metric in self.metrics_to_track:
|
||||
|
||||
Reference in New Issue
Block a user