Fix byol_model_wrapper to function with audio inputs

This commit is contained in:
James Betker 2021-08-05 22:20:22 -06:00
parent f86df53ce0
commit 70dcd1107f

View File

@ -188,19 +188,23 @@ class BYOL(nn.Module):
moving_average_decay=0.99, moving_average_decay=0.99,
use_momentum=True, use_momentum=True,
structural_mlp=False, structural_mlp=False,
positional_dimension=2, # 2 for images, 1 for audio, everything else isn't supported.
perform_augmentation=True,
): ):
super().__init__() super().__init__()
self.online_encoder = NetWrapper(net, projection_size, projection_hidden_size, layer=hidden_layer, self.online_encoder = NetWrapper(net, projection_size, projection_hidden_size, layer=hidden_layer,
use_structural_mlp=structural_mlp) use_structural_mlp=structural_mlp)
augmentations = [ \ self.perform_augmentation = perform_augmentation
RandomApply(augs.ColorJitter(0.8, 0.8, 0.8, 0.2), p=0.8), if self.perform_augmentation:
augs.RandomGrayscale(p=0.2), augmentations = [ \
augs.RandomHorizontalFlip(), RandomApply(augs.ColorJitter(0.8, 0.8, 0.8, 0.2), p=0.8),
RandomApply(filters.GaussianBlur2d((3, 3), (1.5, 1.5)), p=0.1), augs.RandomGrayscale(p=0.2),
augs.RandomResizedCrop((image_size, image_size))] augs.RandomHorizontalFlip(),
self.aug = nn.Sequential(*augmentations) RandomApply(filters.GaussianBlur2d((3, 3), (1.5, 1.5)), p=0.1),
augs.RandomResizedCrop((image_size, image_size))]
self.aug = nn.Sequential(*augmentations)
self.use_momentum = use_momentum self.use_momentum = use_momentum
self.target_encoder = None self.target_encoder = None
self.target_ema_updater = EMA(moving_average_decay) self.target_ema_updater = EMA(moving_average_decay)
@ -212,8 +216,13 @@ class BYOL(nn.Module):
self.to(device) self.to(device)
# send a mock image tensor to instantiate singleton parameters # send a mock image tensor to instantiate singleton parameters
self.forward(torch.randn(2, 3, image_size, image_size, device=device), self.positional_dimension = positional_dimension
torch.randn(2, 3, image_size, image_size, device=device)) if positional_dimension == 2:
self.forward(torch.randn(2, 3, image_size, image_size, device=device),
torch.randn(2, 3, image_size, image_size, device=device))
else:
self.forward(torch.randn(2, 1, 48000, device=device),
torch.randn(2, 1, 48000, device=device))
@singleton('target_encoder') @singleton('target_encoder')
def _get_target_encoder(self): def _get_target_encoder(self):
@ -237,16 +246,17 @@ class BYOL(nn.Module):
return {'target_ema_beta': self.target_ema_updater.beta} return {'target_ema_beta': self.target_ema_updater.beta}
def visual_dbg(self, step, path): def visual_dbg(self, step, path):
torchvision.utils.save_image(self.im1.cpu().float(), os.path.join(path, "%i_image1.png" % (step,))) if self.perform_augmentation and self.positional_dimension == 2:
torchvision.utils.save_image(self.im2.cpu().float(), os.path.join(path, "%i_image2.png" % (step,))) torchvision.utils.save_image(self.im1.cpu().float(), os.path.join(path, "%i_image1.png" % (step,)))
torchvision.utils.save_image(self.im2.cpu().float(), os.path.join(path, "%i_image2.png" % (step,)))
def forward(self, image_one, image_two): def forward(self, image_one, image_two):
image_one = self.aug(image_one.clone()) if self.perform_augmentation:
image_two = self.aug(image_two.clone()) image_one = self.aug(image_one.clone())
image_two = self.aug(image_two.clone())
# Keep copies on hand for visual_dbg. # Keep copies on hand for visual_dbg.
self.im1 = image_one.detach().clone() self.im1 = image_one.detach().clone()
self.im2 = image_two.detach().clone() self.im2 = image_two.detach().clone()
online_proj_one = self.online_encoder(image_one) online_proj_one = self.online_encoder(image_one)
online_proj_two = self.online_encoder(image_two) online_proj_two = self.online_encoder(image_two)
@ -270,4 +280,6 @@ class BYOL(nn.Module):
def register_byol(opt_net, opt): def register_byol(opt_net, opt):
subnet = create_model(opt, opt_net['subnet']) subnet = create_model(opt, opt_net['subnet'])
return BYOL(subnet, opt_net['image_size'], opt_net['hidden_layer'], return BYOL(subnet, opt_net['image_size'], opt_net['hidden_layer'],
structural_mlp=opt_get(opt_net, ['use_structural_mlp'], False)) structural_mlp=opt_get(opt_net, ['use_structural_mlp'], False),
positional_dimension=opt_get(opt_net, ['positional_dims'], 2),
perform_augmentation=opt_get(opt_net, ['aug_enable'], True))