def __init__(
self,
args,
encoder: torch.
nn, # encoder for online network / target network is copy of online network
predictor: torch.
nn, # predictor network comes after online network
optimizer: torch.
optim, # cosine annealing learning rate after warm-up, wo/ restart
**params):
super(BYOL, self).__init__()
self.args = args
self.online_net = encoder.to(args.device)
self.target_net = encoder.to(args.device)
self.predictor = predictor.to(args.device)
self.optimizer = optimizer
self.max_epoch = args.max_epoch
self.batch_size = args.batch_size
self.num_workers = args.num_workers # num_workers for data loader
self.device = args.device # set cuda device
self.writer = SummaryWriter(log_dir=path_summary)
self.resume = args.resume
def __init__(
self,
model: torch.nn.Module,
loss_function: torch.nn,
optimizer: torch.optim,
epochs: int,
model_info: list(),
save_period: int,
savedir: str,
lr_scheduler: torch.optim.lr_scheduler = None,
device: str = None,
):
"""
Args:
model (torch.nn.Module): The model to be trained
loss_function (MultiLoss): The loss function or loss function class
optimizer (torch.optim): torch.optim, i.e., the optimizer class
config (dict): dict of configs
lr_scheduler (torch.optim.lr_scheduler): pytorch lr_scheduler for manipulating the learning rate
seed (int): integer seed to enforce non stochasticity,
device (str): string of the device to be trained on, e.g., "cuda:0"
"""
# Model to device
self.device = torch.device(device)
self.model = model.to(self.device)
self.lr_scheduler = lr_scheduler
self.loss_function = loss_function.to(self.device)
self.optimizer = optimizer
self.epochs = epochs
self.model_info = model_info
self.save_period = save_period
self.start_epoch = 1
self.checkpoint_dir = Path(savedir) / Path(
datetime.today().strftime('%Y-%m-%d'))
self.min_validation_loss = sys.float_info.max # Minimum validation loss achieved, starting with the larges possible number