-
Notifications
You must be signed in to change notification settings - Fork 0
/
inference.py
100 lines (73 loc) · 3.82 KB
/
inference.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
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
import gc
import time
import argparse
import numpy as np
from tqdm import tqdm
import torch
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import DataLoader # noqa
from utils import util
from utils import dataset
from models.unet import UNet
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument('--test_type', type=str, choices=["Temporal", "Spatial"], default='Temporal')
parser.add_argument('--cities', type=list, choices=['ANTWERP', 'BANGKOK', 'BARCELONA', "MOSCOW"], default=['BANGKOK'])
parser.add_argument('--train_year', type=list, choices=[2019, 2020], default=[2020])
parser.add_argument('--val_year', type=list, choices=[2019, 2020], default=[2020])
parser.add_argument('--model', type=str, choices=["UNET"], default="UNET")
parser.add_argument('--batch_size', type=int, default=1)
parser.add_argument('--num_workers', type=int, default=0)
parser.add_argument('--use_mask', type=bool, default=False)
parser.add_argument('--device', type=str, default=torch.device("cuda:0" if torch.cuda.is_available() else "cpu"))
parser.add_argument('--seed', type=int, default=14563)
return parser.parse_args()
def inference(args, path, writer, experiment_name):
_, valid_files, static_map = util.get_files(args)
test_dataset = dataset.TrainDataset(valid_files[0:15], static_map) # using 15/42 files
test_dataloader = DataLoader(dataset=test_dataset, batch_size=args.batch_size, shuffle=False,
num_workers=args.num_workers, collate_fn=dataset.collate_fn, # noqa
sampler=dataset.val_local_sampler(test_dataset), pin_memory=True) # noqa
model = UNet()
checkpoint = torch.load(path)
model.load_state_dict(checkpoint['model_state_dict'])
model.cuda().eval()
criterian = torch.nn.MSELoss()
test_loss = 0.0
running_loss = 0.0
dataset_size = 0
test_predictions = []
test_targets = []
with torch.no_grad():
pbar = tqdm(enumerate(test_dataloader), total=len(test_dataloader), desc="Testing")
for idx, data in pbar:
inputs = data[0].to(args.device, dtype=torch.float)
targets = data[1].to(args.device, dtype=torch.float)
prediction = model(inputs)
pred = prediction[:, :, 1:, 6:-6]
loss = criterian(pred, targets)
dataset_size += inputs.shape[0]
running_loss += loss.item() * inputs.shape[0]
test_loss = running_loss / dataset_size
pbar.set_postfix(test_loss=f"{test_loss:0.4f}")
test_targets.append(targets.cpu().numpy())
test_predictions.append(pred.cpu().numpy())
torch.cuda.empty_cache()
gc.collect()
test_predictions = np.concatenate(test_predictions, axis=0).astype(np.float32)
test_predictions = test_predictions.reshape(-1, 6, 8, 495, 436)
test_predictions = np.moveaxis(test_predictions, source=2, destination=4) # (n, 6, 495, 436, 8)
test_targets = np.concatenate(test_targets, axis=0).astype(np.float32)
test_targets = test_targets.reshape(-1, 6, 8, 495, 436)
test_targets = np.moveaxis(test_targets, source=2, destination=4) # (n, 6, 495, 436, 8)
util.plot(test_predictions, test_targets, experiment_name)
if __name__ == "__main__":
start = time.time()
args = parse_args()
util.seed_everything(args.seed)
experiment_name = ""
writer = SummaryWriter(f"logs/{experiment_name}")
path_to_best = f"checkpoints/{experiment_name}/best_model.pth"
path_to_last = f"checkpoints/{experiment_name}/last_model.pth"
inference(args, path_to_best, writer, experiment_name)
print("Inference_time: ", round((time.time() - start)/60, 2), "(min)")