forked from ownstyledu/Translate-to-Recognize-Networks
-
Notifications
You must be signed in to change notification settings - Fork 0
/
evaluate.py
65 lines (51 loc) · 2.01 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
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
import os
import time
from functools import reduce
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from tensorboardX import SummaryWriter
import data.aligned_conc_dataset as dataset
import util.utils as util
from config.default_config import DefaultConfig
from config.evalute_resnet18_config import EVALUATE_RESNET18_CONFIG
from data import DataProvider
from model.networks import TRecgNet_Upsample_Resiual
cfg = DefaultConfig()
args = {
'resnet18': EVALUATE_RESNET18_CONFIG().args(),
}
# args for different backbones
cfg.parse(args['resnet18'])
os.environ["CUDA_VISIBLE_DEVICES"] = cfg.GPU_IDS
device_ids = torch.cuda.device_count()
print('device_ids:', device_ids)
project_name = reduce(lambda x, y: str(x) + '/' + str(y), os.path.realpath(__file__).split(os.sep)[:-1])
util.mkdir('logs')
val_dataset = dataset.AlignedConcDataset(cfg, data_dir=cfg.DATA_DIR_VAL, transform=transforms.Compose([
dataset.Resize((cfg.LOAD_SIZE, cfg.LOAD_SIZE)),
dataset.CenterCrop((cfg.FINE_SIZE, cfg.FINE_SIZE)),
dataset.ToTensor(),
dataset.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
]))
batch_size_val = cfg.BATCH_SIZE
val_loader = DataProvider(cfg, dataset=val_dataset, batch_size=batch_size_val, shuffle=False)
writer = SummaryWriter(log_dir=cfg.LOG_PATH) # tensorboard
model = TRecgNet_Upsample_Resiual(cfg, writer)
model.set_data_loader(None, val_loader, None)
model.net = nn.DataParallel(model.net).to(model.device)
model.set_log_data(cfg)
def evaluate():
checkpoint_path = os.path.join(cfg.CHECKPOINTS_DIR, cfg.RESUME_PATH)
checkpoint = torch.load(checkpoint_path)
load_epoch = checkpoint['epoch']
model.load_checkpoint(model.net, checkpoint_path, checkpoint, data_para=False)
cfg.START_EPOCH = load_epoch
print('>>> task path is {0}'.format(project_name))
model.evaluate(cfg)
if writer is not None:
writer.close()
if __name__ == '__main__':
start_time = time.time()
evaluate()
print('time consumption: {0} secs', time.time() - start_time)