MLP model
| 73 | |
| 74 | |
| 75 | class MLP(nn.Sequential): |
| 76 | """MLP model""" |
| 77 | |
| 78 | def __init__(self, n_in, n_out, n_hidden=(16, 16, 16), act=None, dropout=0): |
| 79 | if act is None: |
| 80 | act = [ |
| 81 | nn.LeakyReLU(), |
| 82 | ] * (len(n_hidden) + 1) |
| 83 | assert len(act) == len(n_hidden) + 1 |
| 84 | |
| 85 | layer = [] |
| 86 | n_ = [n_in, *n_hidden, n_out] |
| 87 | for i in range(len(n_) - 2): |
| 88 | layer.append(nn.Linear(n_[i], n_[i + 1])) |
| 89 | layer.append(act[i]) |
| 90 | layer.append(nn.Dropout(p=dropout)) |
| 91 | layer.append(nn.Linear(n_[-2], n_[-1])) |
| 92 | super(MLP, self).__init__(*layer) |
| 93 | |
| 94 | |
| 95 | class ConditionalFlowStack(dist.conditional.ConditionalComposeTransformModule): |