Added policy executer file for remotely executing

This commit is contained in:
Victor Mylle
2024-01-15 21:19:33 +00:00
parent 67cc6d4bb9
commit 428e3d9e4b
7 changed files with 323 additions and 110 deletions

View File

@@ -10,7 +10,6 @@ from plotly.subplots import make_subplots
from src.trainers.trainer import Trainer
from tqdm import tqdm
class AutoRegressiveTrainer(Trainer):
def __init__(
self,

View File

@@ -12,6 +12,34 @@ import matplotlib.pyplot as plt
import seaborn as sns
import matplotlib.patches as mpatches
def sample_diffusion(model: DiffusionModel, n: int, inputs: torch.tensor, noise_steps=1000, beta_start=1e-4, beta_end=0.02, ts_length=96):
device = next(model.parameters()).device
beta = torch.linspace(beta_start, beta_end, noise_steps).to(device)
alpha = 1. - beta
alpha_hat = torch.cumprod(alpha, dim=0)
inputs = inputs.repeat(n, 1).to(device)
model.eval()
with torch.no_grad():
x = torch.randn(inputs.shape[0], ts_length).to(device)
for i in reversed(range(1, noise_steps)):
t = (torch.ones(inputs.shape[0]) * i).long().to(device)
predicted_noise = model(x, t, inputs)
_alpha = alpha[t][:, None]
_alpha_hat = alpha_hat[t][:, None]
_beta = beta[t][:, None]
if i > 1:
noise = torch.randn_like(x)
else:
noise = torch.zeros_like(x)
x = 1/torch.sqrt(_alpha) * (x-((1-_alpha) / (torch.sqrt(1 - _alpha_hat))) * predicted_noise) + torch.sqrt(_beta) * noise
return x
class DiffusionTrainer:
def __init__(self, model: nn.Module, data_processor: DataProcessor, device: torch.device):
self.model = model
@@ -50,23 +78,7 @@ class DiffusionTrainer:
def sample(self, model: DiffusionModel, n: int, inputs: torch.tensor):
inputs = inputs.repeat(n, 1).to(self.device)
model.eval()
with torch.no_grad():
x = torch.randn(inputs.shape[0], self.ts_length).to(self.device)
for i in reversed(range(1, self.noise_steps)):
t = (torch.ones(inputs.shape[0]) * i).long().to(self.device)
predicted_noise = model(x, t, inputs)
alpha = self.alpha[t][:, None]
alpha_hat = self.alpha_hat[t][:, None]
beta = self.beta[t][:, None]
if i > 1:
noise = torch.randn_like(x)
else:
noise = torch.zeros_like(x)
x = 1/torch.sqrt(alpha) * (x-((1-alpha) / (torch.sqrt(1 - alpha_hat))) * predicted_noise) + torch.sqrt(beta) * noise
x = sample_diffusion(model, n, inputs, self.noise_steps, self.beta_start, self.beta_end, self.ts_length)
model.train()
return x

View File

@@ -57,6 +57,86 @@ def sample_from_dist(quantiles, output_values):
# Return the mean of the samples
return np.mean(new_samples, axis=1)
def auto_regressive(dataset, model, idx_batch, sequence_length: int = 96):
device = model.device
prev_features, targets = dataset.get_batch(idx_batch)
prev_features = prev_features.to(device)
targets = targets.to(device)
if len(list(prev_features.shape)) == 2:
initial_sequence = prev_features[:, :96]
else:
initial_sequence = prev_features[:, :, 0]
target_full = targets[:, 0].unsqueeze(1) # (batch_size, 1)
with torch.no_grad():
new_predictions_full = model(prev_features) # (batch_size, quantiles)
samples = (
torch.tensor(sample_from_dist( new_predictions_full))
.unsqueeze(1)
.to(device)
) # (batch_size, 1)
predictions_samples = samples
predictions_full = new_predictions_full.unsqueeze(1)
for i in range(sequence_length - 1):
if len(list(prev_features.shape)) == 2:
new_features = torch.cat(
(prev_features[:, 1:96], samples), dim=1
) # (batch_size, 96)
new_features = new_features.float()
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.to(device), other_features.to(device)), dim=1
) # (batch_size, 96 + new_features)
else:
prev_features = new_features
else:
other_features, new_targets = dataset.get_batch_autoregressive(
np.array(idx_batch) + i + 1
) # (batch_size, 1, new_features)
# change the other_features nrv based on the samples
other_features[:, 0, 0] = samples.squeeze(-1)
# make sure on same device
other_features = other_features.to(device)
prev_features = prev_features.to(device)
prev_features = torch.cat(
(prev_features[:, 1:, :], other_features), dim=1
) # (batch_size, 96, new_features)
target_full = torch.cat(
(target_full, new_targets.to(device)), dim=1
) # (batch_size, sequence_length)
with torch.no_grad():
new_predictions_full = model(
prev_features
) # (batch_size, quantiles)
predictions_full = torch.cat(
(predictions_full, new_predictions_full.unsqueeze(1)), dim=1
) # (batch_size, sequence_length, quantiles)
samples = (
torch.tensor(sample_from_dist(new_predictions_full))
.unsqueeze(-1)
.to(device)
) # (batch_size, 1)
predictions_samples = torch.cat((predictions_samples, samples), dim=1)
return (
initial_sequence,
predictions_full,
predictions_samples,
target_full.unsqueeze(-1),
)
class AutoRegressiveQuantileTrainer(AutoRegressiveTrainer):
def __init__(
@@ -273,84 +353,7 @@ class AutoRegressiveQuantileTrainer(AutoRegressiveTrainer):
return fig
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)
if len(list(prev_features.shape)) == 2:
initial_sequence = prev_features[:, :96]
else:
initial_sequence = prev_features[:, :, 0]
target_full = targets[:, 0].unsqueeze(1) # (batch_size, 1)
with torch.no_grad():
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):
if len(list(prev_features.shape)) == 2:
new_features = torch.cat(
(prev_features[:, 1:96], samples), dim=1
) # (batch_size, 96)
new_features = new_features.float()
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.to(self.device), other_features.to(self.device)), dim=1
) # (batch_size, 96 + new_features)
else:
prev_features = new_features
else:
other_features, new_targets = dataset.get_batch_autoregressive(
np.array(idx_batch) + i + 1
) # (batch_size, 1, new_features)
# change the other_features nrv based on the samples
other_features[:, 0, 0] = samples.squeeze(-1)
# make sure on same device
other_features = other_features.to(self.device)
prev_features = prev_features.to(self.device)
prev_features = torch.cat(
(prev_features[:, 1:, :], other_features), dim=1
) # (batch_size, 96, new_features)
target_full = torch.cat(
(target_full, new_targets.to(self.device)), dim=1
) # (batch_size, sequence_length)
with torch.no_grad():
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)
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,
predictions_full,
predictions_samples,
target_full.unsqueeze(-1),
)
return auto_regressive(dataset, self.model, idx_batch, sequence_length)
def plot_quantile_percentages(
self, task, data_loader, train: bool = True, iteration: int = None, full_day: bool = False