def __init__(self, net, image_size, grayscale=True, num_frame_stack=1, batch_size=64, hidden_layer=(- 2), projection_size=256, projection_hidden_size=4096, augment_fn=None, augment_fn2=None, moving_average_decay=0.99, wandb=None, patience=15):
    super().__init__()
    kernel_size = (9, 9)
    kernel_std = np.random.uniform(low=0.1, high=2.0)
    kernel_std = ((kernel_std,) * 2)
    aug_transform = torchvision.transforms.Compose([RandomChoice([enhance.AdjustBrightness(0.4), enhance.AdjustBrightness(0.3), enhance.AdjustBrightness(0.2), enhance.AdjustBrightness(0.1), enhance.AdjustBrightness(0.0)]), RandomChoice([enhance.AdjustContrast(1.0), enhance.AdjustContrast(0.9), enhance.AdjustContrast(0.8), enhance.AdjustContrast(0.7), enhance.AdjustContrast(0.6)]), RandomApply([filters.GaussianBlur2d(kernel_size, kernel_std)], p=0.5)])
    self.augment1 = default(augment_fn, aug_transform)
    self.augment2 = default(augment_fn2, self.augment1)
    self.online_encoder = NetWrapper(net, projection_size, projection_hidden_size, layer=hidden_layer)
    self.target_encoder = None
    self.target_ema_updater = EMA(moving_average_decay)
    self.online_predictor = MLP(projection_size, projection_size, projection_hidden_size)
    self.batch_size = batch_size
    self.device = torch.device(('cuda' if torch.cuda.is_available() else 'cpu'))
    print(f'Device is {self.device.type}')
    self.to(self.device)
    self.wandb = wandb
    self.early_stopper = EarlyStopping(patience=patience, verbose=False, wandb=self.wandb, name='encoder-byol')
    if self.wandb:
        wandb.watch(self.online_encoder, self.target_encoder, self.online_predictor)
    assert grayscale
    nr_channels = num_frame_stack
    self.forward(torch.rand(batch_size, nr_channels, 210, 160, device=self.device))
    self.opt = torch.optim.Adam(self.parameters(), lr=0.0003)
    print(f'Finished Initialization of BYOL with model {self.online_encoder.net.__class__.__name__}')
def __init__(self, net, image_size, grayscale=True, num_frame_stack=1, batch_size=64, hidden_layer=(- 2), projection_size=256, projection_hidden_size=4096, augment_fn=None, augment_fn2=None, moving_average_decay=0.99, wandb=None, patience=15):
    super().__init__()
    kernel_size = (9, 9)
    kernel_std = np.random.uniform(low=0.1, high=2.0)
    kernel_std = ((kernel_std,) * 2)
    aug_transform = torchvision.transforms.Compose([RandomChoice([enhance.AdjustBrightness(0.4), enhance.AdjustBrightness(0.3), enhance.AdjustBrightness(0.2), enhance.AdjustBrightness(0.1), enhance.AdjustBrightness(0.0)]), RandomChoice([enhance.AdjustContrast(1.0), enhance.AdjustContrast(0.9), enhance.AdjustContrast(0.8), enhance.AdjustContrast(0.7), enhance.AdjustContrast(0.6)]), RandomApply([filters.GaussianBlur2d(kernel_size, kernel_std)], p=0.5)])
    self.augment1 = default(augment_fn, aug_transform)
    self.augment2 = default(augment_fn2, self.augment1)
    self.online_encoder = NetWrapper(net, projection_size, projection_hidden_size, layer=hidden_layer)
    self.target_encoder = None
    self.target_ema_updater = EMA(moving_average_decay)
    self.online_predictor = MLP(projection_size, projection_size, projection_hidden_size)
    self.batch_size = batch_size
    self.device = torch.device(('cuda' if torch.cuda.is_available() else 'cpu'))
    print(f'Device is {self.device.type}')
    self.to(self.device)
    self.wandb = wandb
    self.early_stopper = EarlyStopping(patience=patience, verbose=False, wandb=self.wandb, name='encoder-byol')
    if self.wandb:
        wandb.watch(self.online_encoder, self.target_encoder, self.online_predictor)
    assert grayscale
    nr_channels = num_frame_stack
    self.forward(torch.rand(batch_size, nr_channels, 210, 160, device=self.device))
    self.opt = torch.optim.Adam(self.parameters(), lr=0.0003)
    print(f'Finished Initialization of BYOL with model {self.online_encoder.net.__class__.__name__}')