| 12 | |
| 13 | |
| 14 | class HypernetworkModule(torch.nn.Module): |
| 15 | activation_dict = { |
| 16 | "linear": torch.nn.Identity, |
| 17 | "relu": torch.nn.ReLU, |
| 18 | "leakyrelu": torch.nn.LeakyReLU, |
| 19 | "elu": torch.nn.ELU, |
| 20 | "swish": torch.nn.Hardswish, |
| 21 | "tanh": torch.nn.Tanh, |
| 22 | "sigmoid": torch.nn.Sigmoid, |
| 23 | } |
| 24 | activation_dict.update({cls_name.lower(): cls_obj for cls_name, cls_obj in inspect.getmembers(torch.nn.modules.activation) if inspect.isclass(cls_obj) and cls_obj.__module__ == 'torch.nn.modules.activation'}) |
| 25 | |
| 26 | def __init__(self, dim, state_dict=None, layer_structure=None, activation_func=None, weight_init='Normal', |
| 27 | add_layer_norm=False, activate_output=False, dropout_structure=None): |
| 28 | super().__init__() |
| 29 | |
| 30 | self.multiplier = 1.0 |
| 31 | |
| 32 | assert layer_structure is not None, "layer_structure must not be None" |
| 33 | assert layer_structure[0] == 1, "Multiplier Sequence should start with size 1!" |
| 34 | assert layer_structure[-1] == 1, "Multiplier Sequence should end with size 1!" |
| 35 | |
| 36 | linears = [] |
| 37 | for i in range(len(layer_structure) - 1): |
| 38 | |
| 39 | # Add a fully-connected layer |
| 40 | linears.append(torch.nn.Linear(int(dim * layer_structure[i]), int(dim * layer_structure[i+1]))) |
| 41 | |
| 42 | # Add an activation func except last layer |
| 43 | if activation_func == "linear" or activation_func is None or (i >= len(layer_structure) - 2 and not activate_output): |
| 44 | pass |
| 45 | elif activation_func in self.activation_dict: |
| 46 | linears.append(self.activation_dict[activation_func]()) |
| 47 | else: |
| 48 | raise RuntimeError(f'hypernetwork uses an unsupported activation function: {activation_func}') |
| 49 | |
| 50 | # Add layer normalization |
| 51 | if add_layer_norm: |
| 52 | linears.append(torch.nn.LayerNorm(int(dim * layer_structure[i+1]))) |
| 53 | |
| 54 | # Everything should be now parsed into dropout structure, and applied here. |
| 55 | # Since we only have dropouts after layers, dropout structure should start with 0 and end with 0. |
| 56 | if dropout_structure is not None and dropout_structure[i+1] > 0: |
| 57 | assert 0 < dropout_structure[i+1] < 1, "Dropout probability should be 0 or float between 0 and 1!" |
| 58 | linears.append(torch.nn.Dropout(p=dropout_structure[i+1])) |
| 59 | # Code explanation : [1, 2, 1] -> dropout is missing when last_layer_dropout is false. [1, 2, 2, 1] -> [0, 0.3, 0, 0], when its True, [0, 0.3, 0.3, 0]. |
| 60 | |
| 61 | self.linear = torch.nn.Sequential(*linears) |
| 62 | |
| 63 | if state_dict is not None: |
| 64 | self.fix_old_state_dict(state_dict) |
| 65 | self.load_state_dict(state_dict) |
| 66 | else: |
| 67 | for layer in self.linear: |
| 68 | if type(layer) == torch.nn.Linear or type(layer) == torch.nn.LayerNorm: |
| 69 | w, b = layer.weight.data, layer.bias.data |
| 70 | if weight_init == "Normal" or type(layer) == torch.nn.LayerNorm: |
| 71 | normal_(w, mean=0.0, std=0.01) |