pin_memory=True)
# Model
if args.checkpoint is not None and path.exists(args.checkpoint):
model = TorchModel.load_model(args.checkpoint)
else:
network = AnomalyDetector(args.feature_dim)
model = TorchModel(network)
# Training parameters
"""
In the original paper:
lr = 0.01
epsilon = 1e-8
"""
optimizer = torch.optim.Adadelta(network.parameters(),
lr=args.lr_base,
eps=1e-8)
criterion = RegularizedLoss(network, custom_objective)
# Callbacks
tb_writer = SummaryWriter(log_dir=tb_dir)
model.register_callback(
DefaultModelCallback(loss_names=['total'],
visualization_dir=args.exps_dir))
model.register_callback(
TensorBoardCallback(loss_names=['total'], tb_writer=tb_writer))
# Training
model.fit(train_iter=train_iter,
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
cudnn.benchmark = True # enable cudnn tune
train_loader = FeaturesDatasetWrapper(features_path=args.features_path, annotation_path=args.annotation_path)
train_iter = torch.utils.data.DataLoader(train_loader,
batch_size=args.batch_size,
num_workers=0, # 4, # change this part accordingly
pin_memory=True)
network = AnomalyDetector()
system = pw.System(model=network, device=device)
"""
In the original paper:
lr = 0.01
epsilon = 1e-8
"""
optimizer = torch.optim.Adadelta(network.parameters(), lr=args.lr_base, eps=1e-8)
loss_wrapper = pw.loss_wrappers.GenericPointWiseLossWrapper(RegularizedLoss(network, custom_objective))
system.train(
loss_wrapper,
optimizer,
train_data_loader=train_iter,
callbacks=[pw.training_callbacks.NumberOfEpochsStoppingCriterionCallback(args.end_epoch)]
)
system.save_model_state(path.join(args.exps_dir, 'model.weights'))