85 lines
3.2 KiB
Python
85 lines
3.2 KiB
Python
import numpy as np
|
|
import torch
|
|
|
|
from model.pytorch.dcrnn_model import EncoderModel, DecoderModel
|
|
|
|
|
|
class DCRNNSupervisor:
|
|
def __init__(self, adj_mx, **kwargs):
|
|
self._kwargs = kwargs
|
|
self._data_kwargs = kwargs.get('data')
|
|
self._model_kwargs = kwargs.get('model')
|
|
self._train_kwargs = kwargs.get('train')
|
|
|
|
self.num_nodes = int(self._model_kwargs.get('num_nodes', 1))
|
|
self.input_dim = int(self._model_kwargs.get('input_dim', 1))
|
|
self.seq_len = int(self._model_kwargs.get('seq_len')) # for the encoder
|
|
self.output_dim = int(self._model_kwargs.get('output_dim', 1))
|
|
self.use_curriculum_learning = bool(
|
|
self._model_kwargs.get('use_curriculum_learning', False))
|
|
self.horizon = int(self._model_kwargs.get('horizon', 1)) # for the decoder
|
|
|
|
def train(self, encoder_model: EncoderModel, decoder_model: DecoderModel, **kwargs):
|
|
kwargs.update(self._train_kwargs)
|
|
return self._train(**kwargs)
|
|
|
|
def _train_one_batch(self, inputs, labels, encoder_model: EncoderModel,
|
|
decoder_model: DecoderModel, encoder_optimizer,
|
|
decoder_optimizer, criterion):
|
|
"""
|
|
|
|
:param inputs: shape (seq_len, batch_size, num_sensor, input_dim)
|
|
:param labels: shape (horizon, batch_size, num_sensor, input_dim)
|
|
:param encoder_model:
|
|
:param decoder_model:
|
|
:param encoder_optimizer:
|
|
:param decoder_optimizer:
|
|
:param criterion: minimize this criterion
|
|
:return: loss?
|
|
"""
|
|
|
|
encoder_optimizer.zero_grad()
|
|
decoder_optimizer.zero_grad()
|
|
|
|
batch_size = inputs.size(1)
|
|
|
|
inputs = inputs.view(self.seq_len, batch_size, self.num_nodes * self.input_dim)
|
|
labels = labels.view(self.horizon, batch_size, self.num_nodes * self.output_dim)
|
|
|
|
loss = 0
|
|
|
|
encoder_hidden_state = None
|
|
for t in range(self.seq_len):
|
|
_, encoder_hidden_state = encoder_model.forward(inputs[t], encoder_hidden_state)
|
|
|
|
go_symbol = torch.zeros((batch_size, self.num_nodes * self.output_dim))
|
|
|
|
decoder_hidden_state = encoder_hidden_state
|
|
decoder_input = go_symbol
|
|
|
|
for t in range(self.horizon):
|
|
decoder_output, decoder_hidden_state = decoder_model.forward(decoder_input,
|
|
decoder_hidden_state)
|
|
decoder_input = decoder_output
|
|
|
|
if self.use_curriculum_learning: # todo check for is_training (pytorch way?)
|
|
c = np.random.uniform(0, 1)
|
|
if c < self._compute_sampling_threshold():
|
|
decoder_input = labels[t]
|
|
|
|
loss += criterion(decoder_output, labels[t])
|
|
|
|
loss.backward()
|
|
encoder_optimizer.step()
|
|
decoder_optimizer.step()
|
|
return loss.item()
|
|
|
|
def _train(self, encoder_model: EncoderModel, decoder_model: DecoderModel, base_lr, epoch,
|
|
steps, patience=50, epochs=100,
|
|
min_learning_rate=2e-6, lr_decay_ratio=0.1, save_model=1,
|
|
test_every_n_epochs=10):
|
|
pass
|
|
|
|
def _compute_sampling_threshold(self):
|
|
return 1.0 # todo
|