| from typing import Any |
|
|
| import torch |
| from torch import nn |
|
|
| from .base_postprocessor import BasePostprocessor |
|
|
|
|
| class DropoutPostProcessor(BasePostprocessor): |
| def __init__(self, config): |
| self.config = config |
| self.args = config.postprocessor.postprocessor_args |
| self.dropout_times = self.args.dropout_times |
|
|
| @torch.no_grad() |
| def postprocess(self, net: nn.Module, data: Any): |
| logits_list = [net.forward(data) for i in range(self.dropout_times)] |
| logits_mean = torch.zeros_like(logits_list[0], dtype=torch.float32) |
| for i in range(self.dropout_times): |
| logits_mean += logits_list[i] |
| logits_mean /= self.dropout_times |
| score = torch.softmax(logits_mean, dim=1) |
| conf, pred = torch.max(score, dim=1) |
| return pred, conf |
|
|
