| 101 | # modulatory feed forward |
| 102 | |
| 103 | class Modulator(nn.Module): |
| 104 | def __init__(self, dim_in, dim_hidden, num_layers): |
| 105 | super().__init__() |
| 106 | self.layers = nn.ModuleList([]) |
| 107 | |
| 108 | for ind in range(num_layers): |
| 109 | is_first = ind == 0 |
| 110 | dim = dim_in if is_first else (dim_hidden + dim_in) |
| 111 | |
| 112 | self.layers.append(nn.Sequential( |
| 113 | nn.Linear(dim, dim_hidden), |
| 114 | nn.ReLU() |
| 115 | )) |
| 116 | |
| 117 | def forward(self, z): |
| 118 | x = z |
| 119 | hiddens = [] |
| 120 | |
| 121 | for layer in self.layers: |
| 122 | x = layer(x) |
| 123 | hiddens.append(x) |
| 124 | x = torch.cat((x, z)) |
| 125 | |
| 126 | return tuple(hiddens) |
| 127 | |
| 128 | # wrapper |
| 129 | |