-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathmake_square_cam.py
91 lines (70 loc) · 3.03 KB
/
make_square_cam.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
import torch
from torch import multiprocessing, cuda, nn
from torch.utils.data import DataLoader
import torch.nn.functional as F
from torch.backends import cudnn
from net.resnet50_cam import CAM
import argparse
import numpy as np
import os
import voc12.dataloader
from misc import torchutils, pyutils
cudnn.enabled = True
def _work(process_id, model, dataset, config):
cam_out_dir = config['cam_out_dir']
cam_square_shape = config['cam_square_shape']
databin = dataset[process_id]
n_gpus = torch.cuda.device_count()
data_loader = DataLoader(databin, shuffle=False,
num_workers=1, pin_memory=False)
with torch.no_grad(), cuda.device(process_id):
model.cuda()
for iter, pack in enumerate(data_loader):
img_name = pack['name'][0]
label = pack['label'][0]
valid_cat = torch.nonzero(label)[:, 0]
raw_outputs = torch.zeros(
(20, cam_square_shape, cam_square_shape)).cuda()
outputs = [model(img[0].cuda(non_blocking=True))
for img in pack['img']]
outputs = torch.sum(torch.stack(
[F.interpolate(torch.unsqueeze(o, 0),
(cam_square_shape, cam_square_shape), mode='bilinear', align_corners=False)[0]
for o in outputs]), 0)
outputs = outputs[valid_cat]
outputs = outputs / (F.adaptive_max_pool2d(outputs, (1, 1)) + 1e-5)
raw_outputs[valid_cat] = outputs
# save cams
np.save(os.path.join(cam_out_dir, img_name + '.npy'),
{"keys": valid_cat,
"raw_outputs": raw_outputs.cpu().numpy()})
if process_id == n_gpus - 1 and iter % (len(databin) // 20) == 0:
print("%d " % ((5*iter+1)//(len(databin) // 20)), end='')
def run(config):
train_list = config['train_list']
voc12_root = config['voc12_root']
model_root = config['model_root']
cam_scales = config['cam_scales']
cam_weights_name = config['cam_weights_name']
model = CAM()
model.load_state_dict(torch.load(os.path.join(
model_root, cam_weights_name), map_location='cpu'), strict=True)
model.eval()
n_gpus = torch.cuda.device_count()
dataset = voc12.dataloader.VOC12ClassificationDatasetMSF(train_list,
voc12_root=voc12_root,
scales=cam_scales)
dataset = torchutils.split_dataset(dataset, n_gpus)
print('[ ', end='')
multiprocessing.spawn(_work, nprocs=n_gpus, args=(
model, dataset, config), join=True)
print(']')
torch.cuda.empty_cache()
if __name__ == '__main__':
parser = argparse.ArgumentParser(
description='Front Door Semantic Segmentation')
parser.add_argument('--config', type=str,
help='YAML config file path', required=True)
args = parser.parse_args()
config = pyutils.parse_config(args.config)
run(config)