| 37 | |
| 38 | |
| 39 | class SimpleFrozenModel(torch.nn.Module): |
| 40 | |
| 41 | def __init__(self, hidden_dim, empty_grad=False): |
| 42 | super(SimpleFrozenModel, self).__init__() |
| 43 | self.linears = torch.nn.ModuleList([torch.nn.Linear(hidden_dim, hidden_dim) for i in range(2)]) |
| 44 | if empty_grad: |
| 45 | self.linear2 = torch.nn.Linear(hidden_dim, hidden_dim) |
| 46 | self.cross_entropy_loss = torch.nn.CrossEntropyLoss() |
| 47 | self.empty_grad = empty_grad |
| 48 | # Freeze first layer |
| 49 | self.linears[0].weight.requires_grad = False |
| 50 | self.linears[0].bias.requires_grad = False |
| 51 | |
| 52 | def custom_state_dict(self, *args, **kwargs): |
| 53 | state_dict = super(SimpleFrozenModel, self).state_dict(*args, **kwargs) |
| 54 | custom = OrderedDict() |
| 55 | for k, v in state_dict.items(): |
| 56 | if 'linears.0.weight' not in k: |
| 57 | custom[k] = v |
| 58 | return custom |
| 59 | |
| 60 | def forward(self, x, y): |
| 61 | if len(self.linears) == 1: |
| 62 | x = self.linears[0](x) |
| 63 | else: |
| 64 | for i, l in enumerate(self.linears): |
| 65 | x = self.linears[i // 2](x) + l(x) |
| 66 | return self.cross_entropy_loss(x, y) |
| 67 | |
| 68 | |
| 69 | class Curriculum_SimpleModel(SimpleModel): |
no outgoing calls