-
Notifications
You must be signed in to change notification settings - Fork 6
/
evaluate.py
45 lines (40 loc) · 1.25 KB
/
evaluate.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
from argparse import ArgumentParser
import os
import torch
from kornia.losses import SSIMLoss
from kornia.metrics import ssim as compute_ssim
import pytorch_lightning as pl
from pytorch_lightning.loggers import TensorBoardLogger
from pytorch_lightning.callbacks import ModelCheckpoint
from .train import NeuralSupersampling
from .data import testloader
from config import (
learning_rate,
history_length,
upsampling_factor,
ssim_window_size,
n_epochs,
enable_amp,
perceptual_loss_weight,
weight_decay,
tensorboard_root,
source_resolution,
target_resolution,
)
def main(args):
pl.seed_everything(42)
neural_supersampling = NeuralSupersampling()
checkpoint_callback = ModelCheckpoint(dirpath=args.checkpoint_dir, every_n_epochs=1)
trainer = pl.Trainer.from_argparse_args(
args,
accelerator="auto",
precision=16 if enable_amp else 32,
callbacks=[checkpoint_callback],
)
trainer.test(neural_supersampling, testloader)
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument("--checkpoint_dir", type=str, default=os.getcwd(), help="where to load checkpoints from")
parser = pl.Trainer.add_argparse_args(parser)
args = parser.parse_args()
main(args)