
This is the official code of high-resolution representations for ImageNet classification. We augment the HRNet with a classification head shown in the figure below. First, the four-resolution feature maps are fed into a bottleneck and the number of output channels are increased to 128, 256, 512, and 1024, respectively. Then, we downsample the high-resolution representations by a 2-strided 3x3 convolution outputting 256 channels and add them to the representations of the second-high-resolution representations. This process is repeated two times to get 1024 channels over the small resolution. Last, we transform 1024 channels to 2048 channels through a 1x1 convolution, followed by a global average pooling operation. The output 2048-dimensional representation is fed into the classifier.

git: https://github.com/HRNet/HRNet-Image-Classification
Install dependencies: pip install -r requirements.txt
python tools/valid.py --cfg experiments/cls_hrnet_w32_sgd_lr5e-2_wd1e-4_bs32_x100.yaml --testModel models/hrnetv2_w32_imagenet_pretrained.pth修改 NUM_CLASSES
修改 yaml , 作者默认 num_cls = 1000 , MODEL 下 ,增加 NUM_CLASSES
MODEL:
NAME: cls_hrnet
NUM_CLASSES: 2
修改 lib/models/cls_hrnet.py #308
# HighResolutionNet
self.classifier = nn.Linear(2048, cfg['MODEL']['NUM_CLASSES'])
修改 lib/core/function.py accuracy topk
prec1, prec5 = accuracy(output, target, (1, 2))
修改 lib/core/evaluate.py accuracy
#view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.
# 增加 contiguous
correct_k = correct[:k].contiguous().view(-1).float().sum(0, keepdim=True)
准备数据集
├─ datasets
├─ images
│ ├─ train
│ │ ├─ a
│ │ ├─ b
│ ├─ val
│ │ ├─ a
│ │ ├─ b
train
python tools/train.py --cfg experiments/cls_hrnet_w32_sgd_lr5e-2_wd1e-4_bs32_x100.yaml
推理单张图片
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import cv2 as cv
import torch
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
from torch.autograd import Variable
from lib.config import config
from PIL import Image
from lib.models import cls_hrnet
def prediect():
# loda HRNET
config.merge_from_file(
'..\\experiments\\cls_hrnet_w32_sgd_lr5e-2_wd1e-4_bs32_x100.yaml')
config.freeze()
model_file = r"../models/hrnetv2_w32_imagenet_pretrained.pth"
cudnn.benchmark = config.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = config.CUDNN.ENABLED
hrnet = cls_hrnet.get_cls_net(config)
hrnet.load_state_dict(torch.load(model_file))
gpus = list(config.GPUS)
hrnet = torch.nn.DataParallel(hrnet, device_ids=gpus).cuda()
hrnet.eval()
pli_img_path = r'../data/images/000000000139.jpg'
pil_img = Image.open(pli_img_path)
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
input = transforms.Compose([
transforms.Resize(int(config.MODEL.IMAGE_SIZE[0] / 0.875)),
transforms.CenterCrop(config.MODEL.IMAGE_SIZE[0]),
transforms.ToTensor(),
normalize,
])(pil_img)
input = Variable(torch.unsqueeze(input, dim=0).float(), requires_grad=False)
# switch to evaluate mode
cls_pred = 0
with torch.no_grad():
output = hrnet(input)
print(output)
# free image
torch.cuda.empty_cache()
cls_pred = output.argmax(dim=1)
print(cls_pred)
if __name__ == '__main__':
prediect()