| import torch.nn as nn | |
| from torchvision import models | |
| def get_head(out_size, cfg): | |
| """ creates projection head g() from config """ | |
| x = [] | |
| in_size = out_size | |
| for _ in range(cfg.head_layers - 1): | |
| x.append(nn.Linear(in_size, cfg.head_size)) | |
| if cfg.add_bn: | |
| x.append(nn.BatchNorm1d(cfg.head_size)) | |
| x.append(nn.ReLU()) | |
| in_size = cfg.head_size | |
| x.append(nn.Linear(in_size, cfg.emb)) | |
| return nn.Sequential(*x) | |
| def get_model(arch, dataset): | |
| """ creates encoder E() by name and modifies it for dataset """ | |
| model = getattr(models, arch)(pretrained=False) | |
| if dataset != "imagenet": | |
| model.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False) | |
| if dataset == "cifar10" or dataset == "cifar100": | |
| model.maxpool = nn.Identity() | |
| out_size = model.fc.in_features | |
| model.fc = nn.Identity() | |
| return nn.DataParallel(model), out_size | |