Update LFNA

This commit is contained in:
D-X-Y
2021-05-12 20:32:50 +08:00
parent 06f4a1f1cf
commit 0b1ca45c44
8 changed files with 121 additions and 15 deletions

View File

@@ -4,6 +4,8 @@
import copy
import torch
import torch.nn.functional as F
from xlayers import super_core
from xlayers import trunc_normal_
from models.xcore import get_model
@@ -29,13 +31,15 @@ class HyperNet(super_core.SuperModule):
trunc_normal_(self._super_layer_embed, std=0.02)
model_kwargs = dict(
config=dict(model_type="dual_norm_mlp"),
input_dim=layer_embeding + task_embedding,
output_dim=max(self._numel_per_layer),
hidden_dims=[layer_embeding * 4] * 4,
hidden_dims=[layer_embeding * 4] * 3,
act_cls="gelu",
norm_cls="layer_norm_1d",
dropout=0.1,
)
self._generator = get_model(dict(model_type="norm_mlp"), **model_kwargs)
self._generator = get_model(**model_kwargs)
"""
model_kwargs = dict(
input_dim=layer_embeding + task_embedding,
@@ -50,8 +54,12 @@ class HyperNet(super_core.SuperModule):
print("generator: {:}".format(self._generator))
def forward_raw(self, task_embed):
# task_embed = F.normalize(task_embed, dim=-1, p=2)
# layer_embed = F.normalize(self._super_layer_embed, dim=-1, p=2)
layer_embed = self._super_layer_embed
task_embed = task_embed.view(1, -1).expand(self._num_layers, -1)
joint_embed = torch.cat((task_embed, self._super_layer_embed), dim=-1)
joint_embed = torch.cat((task_embed, layer_embed), dim=-1)
weights = self._generator(joint_embed)
if self._return_container:
weights = torch.split(weights, 1)