| from torch import nn |
|
|
|
|
| class LinearClassifier(nn.Module): |
| """Linear layer to train on top of frozen features""" |
|
|
| def __init__( |
| self, |
| dim: int, |
| num_labels: int = 1000, |
| use_dropout_in_head: bool = False, |
| dropout_rate: float = 0.3, |
| large_head: bool = True, |
| use_bn: bool = False, |
| log_softmax: bool = False, |
| ): |
| super(LinearClassifier, self).__init__() |
| self.num_labels = num_labels |
| self.large_head = large_head |
| self.use_bn = use_bn |
| self.log_softmax = log_softmax |
| self.use_dropout_in_head = use_dropout_in_head |
|
|
| if self.use_dropout_in_head: |
| self.dropout = nn.Dropout(dropout_rate) |
| if self.use_bn: |
| self.bn = nn.BatchNorm1d(dim) |
|
|
| if self.large_head: |
| self.linear = nn.Linear(dim, 128) |
| self.linear.weight.data.normal_(mean=0.0, std=0.01) |
| self.linear.bias.data.zero_() |
| self.relu = nn.ReLU() |
|
|
| self.dropout2 = nn.Dropout(dropout_rate) |
| if self.use_bn: |
| self.bn2 = nn.BatchNorm1d(128) |
|
|
| self.linear2 = nn.Linear(128, num_labels) |
| self.linear2.weight.data.normal_(mean=0.0, std=0.01) |
| self.linear2.bias.data.zero_() |
| else: |
| self.linear = nn.Linear(dim, num_labels) |
| self.linear.weight.data.normal_(mean=0.0, std=0.01) |
| self.linear.bias.data.zero_() |
|
|
| def forward(self, x): |
| |
| x = x.view(x.size(0), -1) |
| |
| if self.use_dropout_in_head: |
| x = self.dropout(x) |
| if self.use_bn: |
| x = self.bn(x) |
| |
| x = self.linear(x) |
| |
| if self.large_head: |
| x = self.relu(x) |
| x = self.dropout2(x) |
| if self.use_bn: |
| x = self.bn2(x) |
| |
| x = self.linear2(x) |
| |
| if self.log_softmax: |
| return nn.LogSoftmax(dim=1)(x) |
| else: |
| return x |
|
|
