graph_net_traffic_data/full_model.py at main · NilsHasNoGithub/graph_net_traffic_data · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
from typing import List, Union, Set

import torch
from torch import nn, Tensor
from gnn_model import IntersectionGNN
from roadnet_graph import RoadnetGraph
from vae_net import VariationalEncoderLayer, VariationalLayer, VAEEncoderForwardResult, VAECategoricalDistr, \
    VAELogNormalDistr, VAEDecoderForwardResult, VAEDistr
from dataclasses import dataclass
from typing import Any, Optional
from torch.distributions.normal import Normal
from torch_geometric.data import Batch, Data

@dataclass
class GNNVAEModelState:
    state_dict: Any
    n_features: int
    adj_list: List[List[int]]
    n_out: int
    n_hidden: int
    decoder_distr: VAEDistr

@dataclass
class GNNVAEForwardResult:
    x: Tensor
    kl_div: Tensor
    params_encoder: List[Tensor]
    params_decoder: List[Tensor]


class GNNVAEModel(nn.Module):

    @staticmethod
    def from_model_state(state: GNNVAEModelState) -> "GNNVAEModel":
        model = GNNVAEModel(state.n_features, state.adj_list, n_hidden=state.n_hidden, n_out=state.n_out, decoder_distr=state.decoder_distr)
        model.load_state_dict(state.state_dict)
        return model

    def __init__(self, n_features: int, adj_list: List[List[int]],n_out:int=None, n_hidden: Optional[int]=None, decoder_distr: Optional[VAEDistr]=None):
        """

        :param n_features:
        :param adj_list:
        :param n_out: in case output is not the same as input (should be roughly same size though)
        :param n_hidden: will be automatically chosen if unspecified
        :param decoder_distr: default is lognormal distribution
        """

        nn.Module.__init__(self)

        n_intersetions = len(adj_list)

        edges = []
        for i_from, tos in enumerate(adj_list):
            for i_to in tos:
                edges.append([i_from, i_to])

        self._edges = torch.tensor(edges).transpose(0,1)

        if n_out is None:
            n_out = n_features

        # sizes = [n_features, int(n_features * (5 / 6)), int(n_features * (2 / 3)), int(n_features * (1 / 2))]
        # sizes = [n_features] * 4
        sizes = [n_features] * 3

        if n_hidden is None:
            n_hidden = sizes[-1]

        if decoder_distr is None:
            decoder_distr = VAECategoricalDistr(30)

        self._n_hidden = n_hidden
        self._n_out = n_out
        self._n_features = n_features
        self._adj_list = adj_list


        self._gnn_encoder = IntersectionGNN(sizes, adj_list)
        self._variational_encoder = VariationalEncoderLayer(sizes[-1], n_hidden)
        # self._variational_encoder = VariationalEncoderLayer(sizes[-1], n_out)
        self._gnn_decoder = IntersectionGNN(list(reversed(sizes)), adj_list)
        # Change to not sampling
        self._variational_decoder = VariationalLayer(n_features, n_out, distr=decoder_distr)

        # self._variational_decoder = VariationalLayer(n_features, n_out, distr=VAELogNormalDistr())

    def get_model_state(self) -> GNNVAEModelState:
        return GNNVAEModelState(
            self.state_dict(),
            self._n_features,
            self._adj_list,
            self._n_out,
            self._n_hidden,
            self._variational_decoder.distr
        )

    def distr(self) -> VAEDistr:
        return self._variational_decoder.distr

    def sample(self):

       x = self._variational_encoder.random_output([len(self._adj_list),1, self._n_hidden])

       x = self._gnn_decoder(x, self._edges)

       x: VAEDecoderForwardResult = self._variational_decoder(x)

       return VAEEncoderForwardResult(x.x, torch.tensor(0.0), x.params)


    def forward(self, x: Tensor):
        """

        :param x:
        :param calc_kl_div: Calculate kl div loss and return it if desired
        :return:
        """
        assert x.dim() == 3

        # if isinstance(x, Tensor):
        #     batch_x = self._tensor_to_batch(x)
        # else:
        #     batch_x = x


        x = self._gnn_encoder(x, self._edges)


        encoder_result: VAEEncoderForwardResult = self._variational_encoder(x)

        x = self._gnn_decoder(encoder_result.x, self._edges)
        # x = encoder_result.x

        decoder_result: VAEDecoderForwardResult = self._variational_decoder(x)

        return GNNVAEForwardResult(decoder_result.x, encoder_result.kl_div, encoder_result.params, decoder_result.params)