| 107 | |
| 108 | @persistence.persistent_class |
| 109 | class FullyConnectedLayer(torch.nn.Module): |
| 110 | def __init__(self, |
| 111 | in_features, # Number of input features. |
| 112 | out_features, # Number of output features. |
| 113 | bias = True, # Apply additive bias before the activation function? |
| 114 | activation = 'linear', # Activation function: 'relu', 'lrelu', etc. |
| 115 | lr_multiplier = 1, # Learning rate multiplier. |
| 116 | bias_init = 0, # Initial value for the additive bias. |
| 117 | ): |
| 118 | super().__init__() |
| 119 | self.activation = activation |
| 120 | self.weight = torch.nn.Parameter(torch.randn([out_features, in_features]) / lr_multiplier) |
| 121 | self.bias = torch.nn.Parameter(torch.full([out_features], float(bias_init))) if bias else None |
| 122 | self.weight_gain = lr_multiplier / np.sqrt(in_features) |
| 123 | self.bias_gain = lr_multiplier |
| 124 | |
| 125 | def forward(self, x): |
| 126 | w = self.weight.to(x.dtype) * self.weight_gain |
| 127 | b = self.bias |
| 128 | if b is not None: |
| 129 | b = b.to(x.dtype) |
| 130 | if self.bias_gain != 1: |
| 131 | b = b * self.bias_gain |
| 132 | |
| 133 | if self.activation == 'linear' and b is not None: |
| 134 | x = torch.addmm(b.unsqueeze(0), x, w.t()) |
| 135 | else: |
| 136 | x = x.matmul(w.t()) |
| 137 | x = bias_act.bias_act(x, b, act=self.activation) |
| 138 | return x |
| 139 | |
| 140 | #---------------------------------------------------------------------------- |
| 141 | |