From 226094bb7af1c38f4a5847ffb97cddeebcabd29f Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Wed, 11 Dec 2019 21:51:04 +0100 Subject: [PATCH] works on gpu --- pts/model/forecast.py | 7 +++---- pts/model/forecast_generator.py | 6 +++--- pts/model/predictor.py | 17 +++++++++-------- 3 files changed, 15 insertions(+), 15 deletions(-) diff --git a/pts/model/forecast.py b/pts/model/forecast.py index fc1dc7b..79c40ff 100644 --- a/pts/model/forecast.py +++ b/pts/model/forecast.py @@ -215,8 +215,7 @@ class SampleForecast(Forecast): len(np.shape(samples)) == 2 or len(np.shape(samples)) == 3 ), "samples should be a 2-dimensional or 3-dimensional array. Dimensions found: {}".format( len(np.shape(samples))) - self.samples = samples if (isinstance( - samples, np.ndarray)) else samples.numpy() + self.samples = samples if (isinstance(samples, np.ndarray)) else samples.cpu().numpy() self._sorted_samples_value = None self._mean = None self._dim = None @@ -505,7 +504,7 @@ class DistributionForecast(Forecast): if self._mean is not None: return self._mean else: - self._mean = self.distribution.mean.numpy() + self._mean = self.distribution.mean.cpu().numpy() return self._mean @property @@ -517,7 +516,7 @@ class DistributionForecast(Forecast): def quantile(self, level): level = Quantile.parse(level).value - q = self.distribution.icdf(torch.tensor([level])).numpy() + q = self.distribution.icdf(torch.tensor([level])).cpu().numpy() return q def to_sample_forecast(self, num_samples: int = 200) -> SampleForecast: diff --git a/pts/model/forecast_generator.py b/pts/model/forecast_generator.py index 32fad8f..94f2e11 100644 --- a/pts/model/forecast_generator.py +++ b/pts/model/forecast_generator.py @@ -95,7 +95,7 @@ class QuantileForecastGenerator(ForecastGenerator): num_samples: Optional[int], **kwargs) -> Iterator[Forecast]: for batch in inference_data_loader: inputs = [batch[k] for k in input_names] - outputs = prediction_net(*inputs).numpy() + outputs = prediction_net(*inputs).cpu().numpy() if output_transform is not None: outputs = output_transform(batch, outputs) @@ -120,14 +120,14 @@ class SampleForecastGenerator(ForecastGenerator): num_samples: Optional[int], **kwargs) -> Iterator[Forecast]: for batch in inference_data_loader: inputs = [batch[k] for k in input_names] - outputs = prediction_net(*inputs).numpy() + outputs = prediction_net(*inputs).cpu().numpy() if output_transform is not None: outputs = output_transform(batch, outputs) if num_samples: num_collected_samples = outputs[0].shape[0] collected_samples = [outputs] while num_collected_samples < num_samples: - outputs = prediction_net(*inputs).numpy() + outputs = prediction_net(*inputs).cpu().numpy() if output_transform is not None: outputs = output_transform(batch, outputs) collected_samples.append(outputs) diff --git a/pts/model/predictor.py b/pts/model/predictor.py index 7561eae..e7a27a3 100644 --- a/pts/model/predictor.py +++ b/pts/model/predictor.py @@ -59,11 +59,12 @@ class PTSPredictor(Predictor): device=self.device, dtype=self.dtype, ) - yield from self.forecast_generator( - inference_data_loader=inference_data_loader, - prediction_net=self.prediction_net, - input_names=self.input_names, - freq=self.freq, - output_transform=self.output_transform, - num_samples=num_samples, - ) + with torch.no_grad(): + yield from self.forecast_generator( + inference_data_loader=inference_data_loader, + prediction_net=self.prediction_net, + input_names=self.input_names, + freq=self.freq, + output_transform=self.output_transform, + num_samples=num_samples, + )