def get_learner(config):
resnet = torchvision.models.resnet18(pretrained=True)
learner = BYOL(resnet, image_size=32, hidden_layer='avgpool')
opt = torch.optim.Adam(learner.parameters(), lr=config.lr)
learner = learner.cuda()
return learner
import torchvision
from torchvision import models
import torchvision.transforms as transforms
import torch.nn as nn
from torch.autograd import Variable
use_gpu = torch.cuda.is_available()
resnet = models.resnet50(pretrained=True)
learner = BYOL(resnet, image_size=256
# ,
# hidden_layer = 'avgpool'
)
learner = learner.cuda()
# if torch.cuda.is_available:
# learner=nn.DataParallel(learner,device_ids=[0,1,2]) # multi-GPU
opt = torch.optim.Adam(learner.parameters(), lr=3e-4)
transform = transforms.ToTensor()
trainset = torchvision.datasets.CIFAR10(root='./data',
train=True,
download=True,
transform=transform)
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=64,
shuffle=True,
num_workers=2)
# def sample_unlabelled_images():