-
Notifications
You must be signed in to change notification settings - Fork 5
/
eval.py
executable file
·39 lines (25 loc) · 1.1 KB
/
eval.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
import torch
import argparse
from NeuFlow.neuflow import NeuFlow
from data_utils.evaluate import validate_things, validate_sintel
def get_args_parser():
parser = argparse.ArgumentParser()
parser.add_argument('--resume', default=None, type=str,
help='resume from pretrain model for finetuing or resume from terminated training')
return parser
def main(args):
torch.backends.cudnn.benchmark = True
device = torch.device('cuda')
model = NeuFlow().to(device)
checkpoint = torch.load(args.resume, map_location='cuda')
model.load_state_dict(checkpoint['model'], strict=True)
num_params = sum(p.numel() for p in model.parameters())
print('Number of params:', num_params)
validate_things(model, dstype='frames_cleanpass', validate_subset=False, max_val_flow=400)
validate_things(model, dstype='frames_finalpass', validate_subset=False, max_val_flow=400)
validate_sintel(model, dstype='clean')
validate_sintel(model, dstype='final')
if __name__ == '__main__':
parser = get_args_parser()
args = parser.parse_args()
main(args)