(self, in_dims=80)
| 80 | |
| 81 | class DiffNet(nn.Module): |
| 82 | def __init__(self, in_dims=80): |
| 83 | super().__init__() |
| 84 | self.params = params = AttrDict( |
| 85 | # Model params |
| 86 | encoder_hidden=hparams['hidden_size'], |
| 87 | residual_layers=hparams['residual_layers'], |
| 88 | residual_channels=hparams['residual_channels'], |
| 89 | dilation_cycle_length=hparams['dilation_cycle_length'], |
| 90 | ) |
| 91 | self.input_projection = Conv1d(in_dims, params.residual_channels, 1) |
| 92 | self.diffusion_embedding = SinusoidalPosEmb(params.residual_channels) |
| 93 | dim = params.residual_channels |
| 94 | self.mlp = nn.Sequential( |
| 95 | nn.Linear(dim, dim * 4), |
| 96 | Mish(), |
| 97 | nn.Linear(dim * 4, dim) |
| 98 | ) |
| 99 | self.residual_layers = nn.ModuleList([ |
| 100 | ResidualBlock(params.encoder_hidden, params.residual_channels, 2 ** (i % params.dilation_cycle_length)) |
| 101 | for i in range(params.residual_layers) |
| 102 | ]) |
| 103 | self.skip_projection = Conv1d(params.residual_channels, params.residual_channels, 1) |
| 104 | self.output_projection = Conv1d(params.residual_channels, in_dims, 1) |
| 105 | nn.init.zeros_(self.output_projection.weight) |
| 106 | |
| 107 | def forward(self, spec, diffusion_step, cond): |
| 108 | """ |
no test coverage detected