| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torchvision.models.resnet import resnet50, resnet18 |
|
|
|
|
| class Model(nn.Module): |
| def __init__(self, feature_dim=128, dataset='cifar10', arch='resnet50'): |
| super(Model, self).__init__() |
|
|
| self.f = [] |
| if arch == 'resnet18': |
| temp_model = resnet18().named_children() |
| embedding_size = 512 |
| elif arch == 'resnet50': |
| temp_model = resnet50().named_children() |
| embedding_size = 2048 |
| else: |
| raise NotImplementedError |
| |
| for name, module in temp_model: |
| if name == 'conv1': |
| module = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False) |
| if dataset == 'cifar10' or dataset == 'cifar100': |
| if not isinstance(module, nn.Linear) and not isinstance(module, nn.MaxPool2d): |
| self.f.append(module) |
| elif dataset == 'tiny_imagenet' or dataset == 'stl10': |
| if not isinstance(module, nn.Linear): |
| self.f.append(module) |
| |
| self.f = nn.Sequential(*self.f) |
| |
| self.g = nn.Sequential(nn.Linear(embedding_size, 512, bias=False), nn.BatchNorm1d(512), |
| nn.ReLU(inplace=True), nn.Linear(512, feature_dim, bias=True)) |
|
|
| def forward(self, x): |
| x = self.f(x) |
| feature = torch.flatten(x, start_dim=1) |
| out = self.g(feature) |
| return F.normalize(feature, dim=-1), F.normalize(out, dim=-1) |
|
|
