Added policy executer file for remotely executing
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user