| from typing import Any |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from .base_postprocessor import BasePostprocessor |
|
|
|
|
| class MCMPostprocessor(BasePostprocessor): |
| def __init__(self, config): |
| super(MCMPostprocessor, self).__init__(config) |
| self.args = self.config.postprocessor.postprocessor_args |
| self.tau = self.args.tau |
|
|
| @torch.no_grad() |
| def postprocess(self, net: nn.Module, data: Any): |
| output = net(data) |
| score = torch.softmax(output / self.tau, dim=1) |
| conf, pred = torch.max(score, dim=1) |
| return pred, conf |
|
|
| def set_hyperparam(self, hyperparam: list): |
| self.tau = hyperparam[0] |
|
|
| def get_hyperparam(self): |
| return self.tau |
|
|
