Linear classifier on top of frozen features.
| 153 | |
| 154 | |
| 155 | class LinearClassifier(nn.Module): |
| 156 | """Linear classifier on top of frozen features.""" |
| 157 | |
| 158 | def __init__(self, out_dim: int, use_n_blocks: int, use_avgpool: bool, num_classes: int = 1000): |
| 159 | super().__init__() |
| 160 | self.out_dim = out_dim |
| 161 | self.use_n_blocks = use_n_blocks |
| 162 | self.use_avgpool = use_avgpool |
| 163 | self.num_classes = num_classes |
| 164 | self.linear = nn.Linear(out_dim, num_classes) |
| 165 | self.linear.weight.data.normal_(mean=0.0, std=0.01) |
| 166 | self.linear.bias.data.zero_() |
| 167 | |
| 168 | def forward(self, x_tokens_list): |
| 169 | output = create_linear_input(x_tokens_list, self.use_n_blocks, self.use_avgpool) |
| 170 | return self.linear(output) |
| 171 | |
| 172 | |
| 173 | class AllClassifiers(nn.Module): |
no outgoing calls