def init_encoder(self):
dummy_batch = torch.zeros((2, self.hparams.image_channels, self.hparams.image_height,
self.hparams.image_height))
encoder_name = self.hparams.encoder
if encoder_name == 'amdim_encoder':
encoder = AMDIMEncoder(
dummy_batch,
num_channels=self.hparams.image_channels,
encoder_feature_dim=self.hparams.encoder_feature_dim,
embedding_fx_dim=self.hparams.embedding_fx_dim,
conv_block_depth=self.hparams.conv_block_depth,
encoder_size=self.hparams.image_height,
use_bn=self.hparams.use_bn
)
encoder.init_weights()
return encoder
else:
return torchvision_ssl_encoder(encoder_name, return_all_feature_maps=True)
def __init__(self, hparams):
super().__init__()
self.hparams = hparams
dummy_batch = torch.zeros(
(2, 3, hparams.image_height, hparams.image_height))
self.encoder = AMDIMEncoder(dummy_batch,
num_channels=3,
ndf=hparams.ndf,
n_rkhs=hparams.n_rkhs,
n_depth=hparams.n_depth,
encoder_size=hparams.image_height,
use_bn=hparams.use_bn)
self.encoder.init_weights()
# the loss has learnable parameters
self.nce_loss = AMDIMLossNCE(tclip=self.hparams.tclip)
self.tng_split = None
self.val_split = None