mirror of
https://github.com/open-mmlab/mmpretrain.git
synced 2025-06-03 14:59:18 +08:00
9547e7b7a5
* add model inference on single image * rm --eval * revise doc * add inference tool and demo * fix linting * rename inference_image to inference_model * infer pred_label and pred_score * fix linting * add docstr for inference * add remove_keys * add doc for inference * dump results rather than outputs * add class_names * add related infer scripts * add demo image and the first part of colab tutorial * conduct evaluation in dataset * return lst in simple_test * compuate topk accuracy with numpy * return outputs in test api * merge inference and evaluation tool * fix typo * rm gt_labels in test conifg * get gt_labels during evaluation * sperate the ipython notebook to another PR * return tensor for onnx_export * detach var in simple_test * rm inference script * rm inference script * construct data dict to replace LoadImage * print first predicted result if args.out is None * modify test_pipeline in inference * refactor class_names of imagenet * set class_to_idx as a property in base dataset * output pred_class during inference * remove unused docstr
57 lines
1.7 KiB
Python
57 lines
1.7 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from mmcv.cnn import normal_init
|
|
|
|
from ..builder import HEADS
|
|
from .cls_head import ClsHead
|
|
|
|
|
|
@HEADS.register_module()
|
|
class LinearClsHead(ClsHead):
|
|
"""Linear classifier head.
|
|
|
|
Args:
|
|
num_classes (int): Number of categories excluding the background
|
|
category.
|
|
in_channels (int): Number of channels in the input feature map.
|
|
loss (dict): Config of classification loss.
|
|
""" # noqa: W605
|
|
|
|
def __init__(self,
|
|
num_classes,
|
|
in_channels,
|
|
loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
|
|
topk=(1, )):
|
|
super(LinearClsHead, self).__init__(loss=loss, topk=topk)
|
|
self.in_channels = in_channels
|
|
self.num_classes = num_classes
|
|
|
|
if self.num_classes <= 0:
|
|
raise ValueError(
|
|
f'num_classes={num_classes} must be a positive integer')
|
|
|
|
self._init_layers()
|
|
|
|
def _init_layers(self):
|
|
self.fc = nn.Linear(self.in_channels, self.num_classes)
|
|
|
|
def init_weights(self):
|
|
normal_init(self.fc, mean=0, std=0.01, bias=0)
|
|
|
|
def simple_test(self, img):
|
|
"""Test without augmentation."""
|
|
cls_score = self.fc(img)
|
|
if isinstance(cls_score, list):
|
|
cls_score = sum(cls_score) / float(len(cls_score))
|
|
pred = F.softmax(cls_score, dim=1) if cls_score is not None else None
|
|
if torch.onnx.is_in_onnx_export():
|
|
return pred
|
|
pred = list(pred.detach().cpu().numpy())
|
|
return pred
|
|
|
|
def forward_train(self, x, gt_label):
|
|
cls_score = self.fc(x)
|
|
losses = self.loss(cls_score, gt_label)
|
|
return losses
|