| from torch import nn |
|
|
| from ...nets.attention_model.decoder import Decoder |
| from ...nets.attention_model.embedding import AutoEmbedding |
| from ...nets.attention_model.encoder import GraphAttentionEncoder |
|
|
|
|
| class AttentionModel(nn.Module): |
| r""" |
| The Attention Model from Kool et al.,2019. |
| `Link to paper <https://arxiv.org/abs/1803.08475>`_. |
| |
| For an instance :math:`s`, |
| |
| .. math:: |
| \log(p_\theta(\pi|s)),\pi = \mathrm{AttentionModel}(s) |
| |
| The following is executed: |
| |
| .. math:: |
| \begin{aligned} |
| \pmb{x} &= \mathrm{Embedding}(s) \\ |
| \pmb{h} &= \mathrm{Encoder}(\pmb{x}) \\ |
| \{\log(\pmb{p}_t)\},\pi &= \mathrm{Decoder}(s, \pmb{h}) \\ |
| \log(p_\theta(\pi|s)) &= \sum\nolimits_t\log(\pmb{p}_{t,\pi_t}) |
| \end{aligned} |
| where :math:`\pmb{h}_i` is the node embedding for each node :math:`i` in the graph. |
| |
| In a nutshell, :math:`\mathrm{Embedding}` is a linear projection. |
| The :math:`\mathrm{Encoder}` is a transformer. |
| The :math:`\mathrm{Decoder}` uses (multi-head) attentions. |
| The policy (sequence of action) :math:`\pi` is decoded autoregressively. |
| The log-likelihood :math:`\log(p_\theta(\pi|s))` of the policy is also returned. |
| |
| .. seealso:: |
| The definition of :math:`\mathrm{Embedding}`, :math:`\mathrm{Encoder}`, and |
| :math:`\mathrm{Decoder}` can be found in the |
| :mod:`.embedding`, :mod:`.encoder`, :mod:`.decoder` modules. |
| |
| Args: |
| embedding_dim : the dimension of the embedded inputs |
| hidden_dim : the dimension of the hidden state of the encoder |
| problem : an object defining the state of the problem |
| n_encode_layers: number of encoder layers |
| tanh_clipping : the clipping scale for the decoder |
| Inputs: inputs |
| * **inputs**: problem instance :math:`s`. [batch, graph_size, input_dim] |
| Outputs: ll, pi |
| * **ll**: :math:`\log(p_\theta(\pi|s))` the log-likelihood of the policy. [batch, T] |
| * **pi**: the policy generated :math:`\pi`. [batch, T] |
| """ |
|
|
| def __init__( |
| self, |
| embedding_dim, |
| hidden_dim, |
| problem, |
| n_encode_layers=2, |
| tanh_clipping=10.0, |
| normalization="batch", |
| n_heads=8, |
| ): |
| super(AttentionModel, self).__init__() |
|
|
| self.embedding_dim = embedding_dim |
| self.hidden_dim = hidden_dim |
| self.n_encode_layers = n_encode_layers |
| self.decode_type = None |
| self.n_heads = n_heads |
|
|
| self.problem = problem |
|
|
| self.embedding = AutoEmbedding(problem.NAME, {"embedding_dim": embedding_dim}) |
| step_context_dim = self.embedding.context_dim |
|
|
| self.encoder = GraphAttentionEncoder( |
| n_heads=n_heads, |
| embed_dim=embedding_dim, |
| n_layers=self.n_encode_layers, |
| ) |
| self.decoder = Decoder(embedding_dim, step_context_dim, n_heads, problem, tanh_clipping) |
|
|
| def set_decode_type(self, decode_type): |
| self.decoder.set_decode_type(decode_type) |
|
|
| def forward(self, input): |
| embedding = self.embedding(input) |
| encoded_inputs, _ = self.encoder(embedding) |
| _log_p, pi = self.decoder(input, encoded_inputs) |
| ll = self._calc_log_likelihood(_log_p, pi) |
| return ll, pi |
|
|
| def _calc_log_likelihood(self, _log_p, pi): |
|
|
| |
| log_p = _log_p.gather(2, pi.unsqueeze(-1)).squeeze(-1) |
|
|
| assert (log_p > -1000).data.all(), "Logprobs should not be -inf, check sampling procedure!" |
|
|
| |
| return log_p.sum(1) |
|
|
