| 11 | |
| 12 | |
| 13 | class LinearPred(nn.Module): |
| 14 | |
| 15 | def __init__(self, in_dim, emb_dim, out_dim, num_layer, act='relu'): |
| 16 | super().__init__() |
| 17 | |
| 18 | self.linears = nn.ModuleList() |
| 19 | self.acts = nn.ModuleList() |
| 20 | self.num_layer = num_layer |
| 21 | |
| 22 | for i in range(self.num_layer): |
| 23 | if i == 0: |
| 24 | self.linears.append(nn.Linear(in_dim, emb_dim)) |
| 25 | elif i == self.num_layer - 1: |
| 26 | self.linears.append(nn.Linear(emb_dim, out_dim)) |
| 27 | else: |
| 28 | self.linears.append(nn.Linear(emb_dim, emb_dim)) |
| 29 | if i != self.num_layer - 1: |
| 30 | self.acts.append(obtain_act(act)) |
| 31 | else: |
| 32 | self.acts.append(obtain_act(None)) |
| 33 | |
| 34 | # Initialize parameters |
| 35 | nn.init.xavier_uniform_(self.linears[-1].weight) |
| 36 | if self.linears[-1].bias is not None: |
| 37 | self.linears[-1].bias.data.fill_(0.0) |
| 38 | |
| 39 | def forward(self, x): |
| 40 | |
| 41 | for i in range(self.num_layer): |
| 42 | x = self.acts[i](self.linears[i](x)) |
| 43 | |
| 44 | return x |
| 45 | |
| 46 | |
| 47 | class Encoder(nn.Module): |