File size: 1,970 Bytes
178d33b | 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 | import os.path as osp
from copy import deepcopy
from typing import Any
import torch
from torch import nn
from .base_postprocessor import BasePostprocessor
class EnsemblePostprocessor(BasePostprocessor):
def __init__(self, config):
super(EnsemblePostprocessor, self).__init__(config)
self.config = config
self.postprocess_config = config.postprocessor
self.postprocessor_args = self.postprocess_config.postprocessor_args
assert self.postprocessor_args.network_name == \
self.config.network.name,\
'checkpoint network type and model type do not align!'
# get ensemble args
self.checkpoint_root = self.postprocessor_args.checkpoint_root
# list of trained network checkpoints
self.checkpoints = self.postprocessor_args.checkpoints
# number of networks to esembel
self.num_networks = self.postprocessor_args.num_networks
# get networks
self.checkpoint_dirs = [
osp.join(self.checkpoint_root, path, 'best.ckpt')
for path in self.checkpoints
]
def setup(self, net: nn.Module, id_loader_dict, ood_loader_dict):
self.networks = [deepcopy(net) for i in range(self.num_networks)]
for i in range(self.num_networks):
self.networks[i].load_state_dict(torch.load(
self.checkpoint_dirs[i]),
strict=False)
self.networks[i].eval()
def postprocess(self, net: nn.Module, data: Any):
logits_list = [
self.networks[i](data) for i in range(self.num_networks)
]
logits_mean = torch.zeros_like(logits_list[0], dtype=torch.float32)
for i in range(self.num_networks):
logits_mean += logits_list[i]
logits_mean /= self.num_networks
score = torch.softmax(logits_mean, dim=1)
conf, pred = torch.max(score, dim=1)
return pred, conf
|