import torch
import torch.nn as nn
import torch.nn.functional as F
from libcity.model import loss
from libcity.model.abstract_traffic_state_model import AbstractTrafficStateModel
[docs]def normalize(data, mean, std):
return (data - mean) / std
[docs]def unnormalize(data, mean, std):
return data * std + mean
[docs]class Attr(nn.Module):
def __init__(self, embed_dims, data_feature):
super(Attr, self).__init__()
self.embed_dims = embed_dims
self.data_feature = data_feature
for name, dim_in, dim_out in self.embed_dims:
self.add_module(name + '_em', nn.Embedding(dim_in, dim_out))
[docs] def out_size(self):
sz = 0
for _, _, dim_out in self.embed_dims:
sz += dim_out
# append total distance and holiday
return sz + 2
[docs] def forward(self, batch):
em_list = []
for name, _, _ in self.embed_dims:
embed = getattr(self, name + '_em')
attr_t = batch[name]
attr_t = torch.squeeze(embed(attr_t), dim=1)
em_list.append(attr_t)
dist_mean, dist_std = self.data_feature["dist_mean"], self.data_feature["dist_std"]
dist = normalize(batch["dist"], dist_mean, dist_std)
em_list.append(dist)
em_list.append(batch['holiday'].float())
return torch.cat(em_list, dim=1)
[docs]class ShortSpeed(nn.Module):
def __init__(self, data_feature):
super(ShortSpeed, self).__init__()
self.short_kernel_size = 2
self.short_cnn = nn.Conv1d(3, 4, kernel_size=self.short_kernel_size, stride=1)
self.short_rnn = nn.RNN(input_size=4, hidden_size=16, num_layers=1, batch_first=True)
self.data_feature = data_feature
[docs] def forward(self, batch):
# short-term travel speed features
n_batchs = batch['speeds'].size()[0]
speeds_forward = batch['speeds'].reshape(-1, 4) # (B * L, 4)
speeds_adjacent1 = batch['speeds_relevant1'].reshape(-1, 4)
speeds_adjacent2 = batch['speeds_relevant2'].reshape(-1, 4)
grid_len = batch['grid_len'].reshape(-1, 1) # (B * L, 1)
speeds_forward = torch.unsqueeze(speeds_forward, dim=2) # (B * L, 4, 1)
speeds_adjacent1 = torch.unsqueeze(speeds_adjacent1, dim=2)
speeds_adjacent2 = torch.unsqueeze(speeds_adjacent2, dim=2)
grid_len = torch.unsqueeze(grid_len, dim=2) # (B * L, 1, 1)
grid_len_short = grid_len.expand(speeds_forward.size()[:2] + (grid_len.size()[-1], )) # (B * L, 4, 1)
times_forward = speeds_forward.clone() # (B * L, 4, 1)
times_forward[times_forward==0] = 0.2
times_forward = grid_len_short / times_forward * 3600
times_adjacent1 = speeds_adjacent1.clone()
times_adjacent1[times_adjacent1==0] = 0.2
times_adjacent1 = grid_len_short / times_adjacent1 * 3600
times_adjacent2 = speeds_adjacent2.clone()
times_adjacent2[times_adjacent2==0] = 0.2
times_adjacent2 = grid_len_short / times_adjacent2 * 3600
speeds_forward = normalize(speeds_forward, self.data_feature["speeds_mean"], self.data_feature["speeds_std"])
speeds_adjacent1 = normalize(speeds_adjacent1, self.data_feature["speeds_relevant1_mean"], self.data_feature["speeds_relevant1_std"])
speeds_adjacent2 = normalize(speeds_adjacent2, self.data_feature["speeds_relevant2_mean"], self.data_feature["speeds_relevant2_std"])
grid_len_short = normalize(grid_len_short, self.data_feature["grid_len_mean"], self.data_feature["grid_len_std"])
times_forward = normalize(times_forward, self.data_feature["time_gap_mean"], self.data_feature["time_gap_std"])
times_adjacent1 = normalize(times_adjacent1, self.data_feature["time_gap_mean"], self.data_feature["time_gap_std"])
times_adjacent2 = normalize(times_adjacent2, self.data_feature["time_gap_mean"], self.data_feature["time_gap_std"])
inputs_0 = torch.cat([speeds_forward, grid_len_short, times_forward], dim=2) # (B * L, 4, 3)
inputs_1 = torch.cat([speeds_adjacent1, grid_len_short, times_adjacent1], dim=2) # (B * L, 4, 3)
inputs_2 = torch.cat([speeds_adjacent2, grid_len_short, times_adjacent2], dim=2) # (B * L, 4, 3)
outputs_0 = torch.tanh(self.short_cnn(inputs_0.permute(0, 2, 1)))
outputs_0 = outputs_0.permute(0, 2, 1)
outputs_1 = torch.tanh(self.short_cnn(inputs_1.permute(0, 2, 1)))
outputs_1 = outputs_1.permute(0, 2, 1)
outputs_2 = torch.tanh(self.short_cnn(inputs_2.permute(0, 2, 1)))
outputs_2 = outputs_2.permute(0, 2, 1)
outputs_0, _ = self.short_rnn(outputs_0)
outputs_1, _ = self.short_rnn(outputs_1)
outputs_2, _ = self.short_rnn(outputs_2)
outputs_0 = outputs_0.reshape(n_batchs, -1, 4-self.short_kernel_size+1, 16)
outputs_1 = outputs_1.reshape(n_batchs, -1, 4-self.short_kernel_size+1, 16)
outputs_2 = outputs_2.reshape(n_batchs, -1, 4-self.short_kernel_size+1, 16)
V_short = torch.cat([outputs_0[:, :, -1], outputs_1[:, :, -1], outputs_2[:, :, -1]], dim=2)
return V_short
[docs]class LongSpeed(nn.Module):
def __init__(self, data_feature):
super(LongSpeed, self).__init__()
self.long_kernel_size = 3
self.long_cnn = nn.Conv1d(3, 4, kernel_size=self.long_kernel_size, stride=1)
self.long_rnn = nn.RNN(input_size=4, hidden_size=16, num_layers=1, batch_first=True)
self.data_feature = data_feature
[docs] def forward(self, batch):
n_batchs = batch['speeds_long'].size()[0]
speeds_history = batch['speeds_long'].reshape(-1, 7)
grid_len = batch['grid_len'].reshape(-1, 1)
speeds_history = torch.unsqueeze(speeds_history, dim=2)
grid_len = torch.unsqueeze(grid_len, dim=2)
grid_len_long = grid_len.expand(speeds_history.size()[:2] + (grid_len.size()[-1], ))
times_history = speeds_history.clone()
times_history[times_history==0] = 0.2
times_history = grid_len_long / times_history * 3600
speeds_history = normalize(speeds_history, self.data_feature["speeds_long_mean"], self.data_feature["speeds_long_std"])
grid_len_long = normalize(grid_len_long, self.data_feature["grid_len_mean"], self.data_feature["grid_len_std"])
times_history = normalize(times_history, self.data_feature["time_gap_mean"], self.data_feature["time_gap_std"])
inputs_3 = torch.cat([speeds_history, grid_len_long, times_history], dim=2)
outputs_3 = self.long_cnn(inputs_3.permute(0, 2, 1))
outputs_3 = outputs_3.permute(0, 2, 1)
outputs_3, _ = self.long_rnn(outputs_3)
outputs_3 = outputs_3.reshape(n_batchs, -1, 7 - self.long_kernel_size + 1, 16)
V_long = outputs_3[:, :, -1]
return V_long
[docs]class SpeedLSTM(nn.Module):
def __init__(self, data_feature):
super(SpeedLSTM, self).__init__()
self.shortspeed_net = ShortSpeed(data_feature)
self.longspeed_net = LongSpeed(data_feature)
self.process_speeds = nn.Linear(64, 32)
self.speed_lstm = nn.LSTM(
input_size = 32,
hidden_size = 32,
num_layers = 1,
batch_first = True,
bidirectional = False,
dropout = 0,
)
[docs] def forward(self, batch):
shortspeeds_t = self.shortspeed_net(batch)
longspeeds_t = self.longspeed_net(batch)
whole_t = torch.cat([shortspeeds_t, longspeeds_t], dim=2)
whole_t = self.process_speeds(whole_t)
whole_t = torch.tanh(whole_t)
lens = [batch["current_dis"].shape[1]] * batch["current_dis"].shape[0]
lens = list(map(lambda x: x, lens))
packed_inputs = nn.utils.rnn.pack_padded_sequence(whole_t, lens, batch_first=True)
packed_hiddens, (_, _) = self.speed_lstm(packed_inputs)
speeds_hiddens, _ = nn.utils.rnn.pad_packed_sequence(packed_hiddens, batch_first=True)
return speeds_hiddens
[docs]class Road(nn.Module):
def __init__(self, data_feature):
super(Road, self).__init__()
self.data_feature = data_feature
self.embedding = nn.Embedding(128 * 128, 32)
emb_vectors = data_feature["geo_embedding"]
self.embedding.weight.data.copy_(torch.tensor(emb_vectors))
self.process_coords = nn.Linear(2 + 32, 32)
[docs] def forward(self, batch):
longi_mean, longi_std = self.data_feature["longi_mean"], self.data_feature["longi_std"]
current_longi = normalize(batch["current_longi"], longi_mean, longi_std)
lngs = torch.unsqueeze(current_longi, dim=2)
lati_mean, lati_std = self.data_feature["lati_mean"], self.data_feature["lati_std"]
current_lati = normalize(batch["current_lati"], lati_mean, lati_std)
lats = torch.unsqueeze(current_lati, dim=2)
grid_ids = torch.unsqueeze(batch['current_loc'].long(), dim=2)
grids = torch.squeeze(self.embedding(grid_ids), dim=2)
locs = torch.cat([lngs, lats, grids], dim=2)
locs = self.process_coords(locs)
locs = torch.tanh(locs)
return locs
[docs]class RoadLSTM(nn.Module):
def __init__(self, data_feature):
super(RoadLSTM, self).__init__()
self.Road_net = Road(data_feature)
self.Road_lstm = nn.LSTM(
input_size=32,
hidden_size=32,
num_layers=1,
batch_first=True,
bidirectional=False,
dropout=0,
)
[docs] def forward(self, batch):
Roads_t = self.Road_net(batch)
whole_t = Roads_t
lens = [batch["current_dis"].shape[1]] * batch["current_dis"].shape[0]
lens = list(map(lambda x: x, lens))
packed_inputs = nn.utils.rnn.pack_padded_sequence(whole_t, lens, batch_first=True)
packed_hiddens, (_, _) = self.Road_lstm(packed_inputs)
Roads_hiddens, _ = nn.utils.rnn.pad_packed_sequence(packed_hiddens, batch_first=True)
return Roads_hiddens
[docs]class PredictionBiLSTM(nn.Module):
def __init__(self, embed_dims, data_feature):
super(PredictionBiLSTM, self).__init__()
self.attr_net = Attr(embed_dims, data_feature)
self.speed_lstm = SpeedLSTM(data_feature)
self.road_lstm = RoadLSTM(data_feature)
self.bi_lstm = nn.LSTM(
input_size=self.attr_net.out_size() + 64,
hidden_size=64,
num_layers=2,
batch_first=True,
bidirectional=True,
dropout=0.25,
)
self.lnhiddens = nn.LayerNorm(self.attr_net.out_size() + 64, elementwise_affine=True)
[docs] def forward(self, batch):
speeds_t = self.speed_lstm(batch)
roads_t = self.road_lstm(batch)
attr_t = self.attr_net(batch)
attr_t = torch.unsqueeze(attr_t, dim=1)
expand_attr_t = attr_t.expand(roads_t.size()[:2] + (attr_t.size()[-1], ))
hiddens = torch.cat([expand_attr_t, speeds_t, roads_t], dim=2)
hiddens = self.lnhiddens(hiddens)
lens = [batch["current_longi"].shape[1]] * batch["current_longi"].shape[0]
lens = list(map(lambda x: x, lens))
packed_inputs = nn.utils.rnn.pack_padded_sequence(hiddens, lens, batch_first=True)
packed_hiddens, (_, _) = self.bi_lstm(packed_inputs)
hiddens, lens = nn.utils.rnn.pad_packed_sequence(packed_hiddens, batch_first=True)
return hiddens
[docs]class TTPNet(AbstractTrafficStateModel):
def __init__(self, config, data_feature):
super(TTPNet, self).__init__(config, data_feature)
self.config = config
self.data_feature = data_feature
self.device = config.get('device', torch.device('cpu'))
uid_emb_size = config.get("uid_emb_size", 16)
weekid_emb_size = config.get("weekid_emb_size", 3)
timdid_emb_size = config.get("timdid_emb_size", 8)
uid_size = data_feature.get("uid_size", 13000)
embed_dims = [
('uid', uid_size, uid_emb_size),
('weekid', 7, weekid_emb_size),
('timeid', 96, timdid_emb_size),
]
self.bi_lstm = PredictionBiLSTM(embed_dims, data_feature)
self.input2hid = nn.Linear(128, 128)
self.hid2hid = nn.Linear(128, 64)
self.hid2out = nn.Linear(64, 1)
[docs] def forward(self, batch):
hiddens = self.bi_lstm(batch)
n = hiddens.size()[1]
h_f = []
for i in range(2, n):
h_f_temp = torch.sum(hiddens[:, :i], dim=1)
h_f.append(h_f_temp)
h_f.append(torch.sum(hiddens, dim=1))
h_f = torch.stack(h_f).permute(1, 0, 2)
T_f_hat = self.input2hid(h_f)
T_f_hat = F.relu(T_f_hat)
T_f_hat = self.hid2hid(T_f_hat)
T_f_hat = F.relu(T_f_hat)
T_f_hat = self.hid2out(T_f_hat)
return T_f_hat.squeeze(dim=2)
[docs] def calculate_loss(self, batch):
if self.training:
T_f_hat = self.predict(batch).unsqueeze(dim=2)
T_f = torch.unsqueeze(batch["current_tim"][:, 1:], dim=2)
M_f = torch.unsqueeze(batch["masked_current_tim"][:, 1:], dim=1)
loss_f = torch.bmm(M_f, torch.pow((T_f_hat - T_f) / T_f, 2)) / torch.bmm(M_f, M_f.permute(0, 2, 1))
loss_f = torch.pow(loss_f, 1/2)
return loss_f.mean()
else:
pred = self.predict(batch)
label = batch["time"]
return loss.masked_mape_torch(pred, label)
[docs] def predict(self, batch):
T_f_hat = self.forward(batch)
time_gap_mean, time_gap_std = self.data_feature["time_gap_mean"], self.data_feature["time_gap_std"]
T_f_hat = unnormalize(T_f_hat, time_gap_mean, time_gap_std)
if self.training:
return T_f_hat
else:
return T_f_hat[:, -1:]