| 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 |