def __init__(self, hparams):
super().__init__()
self.hparams = hparams
self.online_evaluator = self.hparams.online_ft
self.dataset = self.get_dataset(hparams.dataset)
# encoder network (Z vectors)
dummy_batch = torch.zeros(
(2, 3, hparams.patch_size, hparams.patch_size))
self.encoder = CPCResNet101(dummy_batch)
# info nce loss
c, h = self.__compute_final_nb_c(hparams.patch_size)
self.info_nce = InfoNCE(num_input_channels=c,
target_dim=64,
embed_scale=0.1)
if self.online_evaluator:
z_dim = c * h * h
num_classes = self.dataset.num_classes
self.non_linear_evaluator = SSLEvaluator(n_input=z_dim,
n_classes=num_classes,
p=0.2,
n_hidden=1024)
def init_encoder(self):
dummy_batch = torch.zeros((2, 3, self.hparams.patch_size, self.hparams.patch_size))
encoder_name = self.hparams.encoder
if encoder_name == 'cpc_encoder':
return CPCResNet101(dummy_batch)
else:
return torchvision_ssl_encoder(encoder_name, return_all_feature_maps=self.hparams.task == 'amdim')