mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix byol_model_wrapper to function with audio inputs

Browse Source
This commit is contained in:
James Betker 2021-08-05 22:20:22 -06:00
parent f86df53ce0
commit 70dcd1107f
1 changed files with 30 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
codes/models/byol/byol_model_wrapper.py
View File

@ -188,12 +188,16 @@ 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)
self.perform_augmentation = perform_augmentation
if self.perform_augmentation:
augmentations = [ \ augmentations = [ \
RandomApply(augs.ColorJitter(0.8, 0.8, 0.8, 0.2), p=0.8), RandomApply(augs.ColorJitter(0.8, 0.8, 0.8, 0.2), p=0.8),
augs.RandomGrayscale(p=0.2), augs.RandomGrayscale(p=0.2),
@ -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.positional_dimension = positional_dimension
if positional_dimension == 2:
self.forward(torch.randn(2, 3, image_size, image_size, device=device), self.forward(torch.randn(2, 3, image_size, image_size, device=device),
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,13 +246,14 @@ 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):
if self.perform_augmentation and self.positional_dimension == 2:
torchvision.utils.save_image(self.im1.cpu().float(), os.path.join(path, "%i_image1.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,))) 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):
if self.perform_augmentation:
image_one = self.aug(image_one.clone()) image_one = self.aug(image_one.clone())
image_two = self.aug(image_two.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()
@ -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))