learning_quantum_complexity/model.py at master · markusschmitt/learning_quantum_complexity · 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
#
# Definition of autoencoder networks
#
# Author: Markus Schmitt
# Date: Feb 2021
#

import torch
from torch import nn

class BottleneckNet(nn.Module):
    """This class defines the autoencoder network

    Initializer arguments:

        * ``num_latent``: Number of latent variables (integer).
        * ``encoder_dim``: Sizes of the encoder layers (list of integers)
        * ``decoder_dim``: Sizes of the decoder layers (list of integers)
    """

    def __init__(self, num_latent, input_dim, encoder_dim=[100,100], decoder_dim=[100,100]):
        """Initializes ``BottleneckNet``.

        Arguments:

            * ``num_latent``: Number of latent variables (integer).
            * ``input_dim``: Dimension of input (integer).
            * ``encoder_dim``: Sizes of the encoder layers (list of integers)
            * ``decoder_dim``: Sizes of the decoder layers (list of integers)
        """

        super().__init__()

        # List of encoder layers
        self.enc = [nn.Linear(input_dim, encoder_dim[0])]
        for k in range(len(encoder_dim)-2):
            self.enc.append(nn.Linear(encoder_dim[k],encoder_dim[k+1]))
        self.enc.append(nn.Linear(encoder_dim[-1], num_latent))

        self.enc = nn.ModuleList(self.enc)

        # List of decoder layers
        self.dec = [nn.Linear(num_latent, decoder_dim[0])]
        for k in range(len(decoder_dim)-2):
            self.dec.append(nn.Linear(decoder_dim[k],decoder_dim[k+1]))
        self.dec.append(nn.Linear(decoder_dim[-1], input_dim))

        self.dec = nn.ModuleList(self.dec)

    def forward(self, x):
        """Network evaluation.

        Arguments:

            * ``x``: Input data (Pytorch tensor of floats)

        Returns: Network output (same dimensions as input), latent values
        """

        # Encoder
        for layer in self.enc:
            x=layer(x)
            x=torch.tanh(x)

        # Store latent values
        latent=x.clone()

        # Decoder
        for layer in self.dec:
            x=layer(x)
            x=torch.tanh(x)

        return x, latent

    def compute_latent(self, x):
        """Compute latent values.

        Evaluate only the encoder part of the network

        Arguments:

            * ``x``: Input data (Pytorch tensor of floats)

        Returns: latent values
        """

        # Encoder
        for layer in self.enc:
            x=layer(x)
            x=torch.tanh(x)
        return x

class NoLatentNet(nn.Module):
    """This class mimics the ''autoencoder'' network without latent variables.

    Initializer arguments:

        * ``input_dim``: Dimension of input (integer).
    """
    def __init__(self, input_dim):
        """Initialize ``NoLatentNet``

        Arguments:

            * ``input_dim``: Dimension of input (integer).
        """

        super().__init__()

        # Parameters W correspond to the constant output.
        self.W = torch.nn.Parameter(torch.randn(input_dim))
        self.W.requires_grad = True

    def forward(self, x):
        """Network evaluation.

        The network output is independent of the input.

        Arguments:

            * ``x``: Input data (Pytorch tensor of floats)

        Returns: Network output (same dimensions as input), latent values (dummy)
        """

        return torch.stack([self.W for i in x]), torch.tensor([0.])