Repository: QLMX/huawei-garbage Branch: master Commit: 830580e940b7 Files: 35 Total size: 137.0 KB Directory structure: gitextract_gzr6tbsm/ ├── README.md ├── args.py ├── build_net.py ├── dataset.py ├── models/ │ ├── Res.py │ ├── __init__.py │ ├── resnet_cbam.py │ └── resnetxt_wsl.py ├── predict/ │ ├── Res.py │ ├── config.json │ ├── customize_service.py │ ├── predict.py │ ├── resnet_cbam.py │ └── resnetxt_wsl.py ├── preprocess.py ├── requirements.txt ├── train.py ├── transform.py └── utils/ ├── __init__.py ├── eval.py ├── logger.py ├── misc.py ├── progress/ │ ├── LICENSE │ ├── MANIFEST.in │ ├── README.rst │ ├── progress/ │ │ ├── __init__.py │ │ ├── bar.py │ │ ├── counter.py │ │ ├── helpers.py │ │ └── spinner.py │ ├── setup.py │ └── test_progress.py ├── radam.py ├── utils.py └── visualize.py ================================================ FILE CONTENTS ================================================ ================================================ FILE: README.md ================================================ ## 华为云垃圾分类挑战杯亚军方案分享 ### 1.代码结构 ``` {repo_root} ├── models //模型文件夹 ├── utils //一些函数包 | ├── eval.py // 求精度 │ ├── misc.py // 模型保存，参数初始化，优化函数选择 │ ├── radam.py │ └── ... ├── args.py //参数配置文件 ├── build_net.py //搭建模型 ├── dataset.py //数据批量加载文件 ├── preprocess.py //数据预处理文件，生成坐标标签 ├── train.py //训练运行文件 ├── transform.py //数据增强文件 ``` ### 2. 环境设置 可以直接通过`pip install -r requirements.txt`安装指定的函数包，python版本为3.6，具体的函数包如下： * pytorch>=1.0.1 * torchvision==0.2.2 * matplotlib>=3.1.0 * numpy>=1.16.4 * scikit-image * pandas * sklearn 注：py3.7训练的话，要修改下面的代码 `if use_cuda: inputs, targets = inputs.cuda(), targets.cuda(async=True) inputs, targets = torch.autograd.Variable(inputs), torch.autograd.Variable(targets)` \#python3.7已经移除了async关键字，而用non_blocking代替。(导致apache-airflow也出了问题) \#cuda() 本身也没有async. 就是把 async=True去掉 if use_cuda: inputs, targets = inputs.cuda(), targets.cuda() inputs, targets = torch.autograd.Variable(inputs), torch.autograd.Variable(targets)` ## 3.运行步骤 1. 建立文件夹data，把garbage_classify全部解压缩到data下 2. 运行preprocess.py，生成训练集和测试集运行文 3. 单张显卡的话，修改arg.py 85行 parser.add_argument('--gpu-id', default='0, 1, 2, 3' 为'--gpu-id', default='0'，同时修改 '--train-batch'，'--test-batch'为适当的数字 4. 运行train.py ### 4.方案思路 [方案讲解](https://mp.weixin.qq.com/s/7GhXMXQkBgH_JVcKMjCejQ) 知乎专栏：[ML与DL成长之路](https://zhuanlan.zhihu.com/ai-growth) 如果复现过程中有bug，麻烦反馈一下，会优化更新。如果对您有帮助记得给个**star** --- **小尾巴** QQ群：AI成长社①：545702197 微信群：添加微信号：Derek_wen8，备注：加群 ================================================ FILE: args.py ================================================ #!usr/bin/env python #-*- coding:utf-8 _*- """ @version: python3.6 @author: QLMX @contact: wenruichn@gmail.com @time: 2019-08-15 14:25 公众号：AI成长社 知乎：https://www.zhihu.com/people/qlmx-61/columns """ import argparse from build_net import model_names # Parse arguments parser = argparse.ArgumentParser(description='PyTorch ImageNet Training') # Datasets parser.add_argument('-train', '--trainroot', default='data/new_shu_label.txt', type=str) #new_shu_label parser.add_argument('-val', '--valroot', default='data/val1.txt', type=str) parser.add_argument('-j', '--workers', default=4, type=int, metavar='N', help='number of data loading workers (default: 4)') # Optimization options parser.add_argument('--epochs', default=200, type=int, metavar='N', help='number of total epochs to run') parser.add_argument('--num-classes', default=43, type=int, metavar='N', help='number of classfication of image') parser.add_argument('--image-size', default=288, type=int, metavar='N', help='the train image size') parser.add_argument('--start-epoch', default=0, type=int, metavar='N', help='manual epoch number (useful on restarts)') parser.add_argument('--train-batch', default=64, type=int, metavar='N', help='train batchsize (default: 256)') parser.add_argument('--test-batch', default=32, type=int, metavar='N', help='test batchsize (default: 200)') parser.add_argument('--optimizer', default='sgd', choices=['sgd', 'rmsprop', 'adam', 'AdaBound', 'radam'], metavar='N', help='optimizer (default=sgd)') parser.add_argument('--lr', '--learning-rate', default=0.001, type=float, metavar='LR', help='initial learning rate，1e-2， 1e-4, 0.001') parser.add_argument('--lr-fc-times', '--lft', default=5, type=int, metavar='LR', help='initial model last layer rate') parser.add_argument('--drop', '--dropout', default=0, type=float, metavar='Dropout', help='Dropout ratio') parser.add_argument('--schedule', type=int, nargs='+', default=[30, 50, 60], help='Decrease learning rate at these epochs.') parser.add_argument('--gamma', type=float, default=0.1, help='LR is multiplied by gamma on schedule.') parser.add_argument('--momentum', default=0.9, type=float, metavar='M', help='momentum') parser.add_argument('--no_nesterov', dest='nesterov', action='store_false', help='do not use Nesterov momentum') parser.add_argument('--alpha', default=0.99, type=float, metavar='M', help='alpha for ') parser.add_argument('--beta1', default=0.9, type=float, metavar='M', help='beta1 for Adam (default: 0.9)') parser.add_argument('--beta2', default=0.999, type=float, metavar='M', help='beta2 for Adam (default: 0.999)') parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float, metavar='W', help='weight decay (default: 1e-4)') parser.add_argument('--final-lr', '--fl', default=1e-3,type=float, metavar='W', help='weight decay (default: 1e-3)') # Checkpoints parser.add_argument('-c', '--checkpoint', default='/data0/search/qlmx/clover/garbage/res_16_288_last1', type=str, metavar='PATH', help='path to save checkpoint (default: checkpoint)') parser.add_argument('--resume', default='', type=str, metavar='PATH', help='path to latest checkpoint (default: none)') # Architecture parser.add_argument('--arch', '-a', metavar='ARCH', default='resnext101_32x16d_wsl', choices=model_names, help='model architecture: ' + ' | '.join(model_names) + ' (default: resnext101_32x8d, pnasnet5large)') parser.add_argument('--depth', type=int, default=29, help='Model depth.') parser.add_argument('--cardinality', type=int, default=32, help='ResNet cardinality (group).') parser.add_argument('--base-width', type=int, default=4, help='ResNet base width.') parser.add_argument('--widen-factor', type=int, default=4, help='Widen factor. 4 -> 64, 8 -> 128, ...') # Miscs parser.add_argument('--manualSeed', type=int, help='manual seed') parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true', help='evaluate model on validation set') parser.add_argument('--pretrained', dest='pretrained', action='store_true', help='use pre-trained model') #Device options parser.add_argument('--gpu-id', default='0, 1, 2, 3', type=str, help='id(s) for CUDA_VISIBLE_DEVICES') args = parser.parse_args() ================================================ FILE: build_net.py ================================================ #!usr/bin/env python #-*- coding:utf-8 _*- """ @version: python3.6 @author: QLMX @contact: wenruichn@gmail.com @time: 2019-08-15 14:30 """ from torch import nn import torchvision.models as models import models as customized_models # Models default_model_names = sorted(name for name in models.__dict__ if name.islower() and not name.startswith("__") and callable(models.__dict__[name])) customized_models_names = sorted(name for name in customized_models.__dict__ if not name.startswith("__") and callable(customized_models.__dict__[name])) for name in customized_models.__dict__: if not name.startswith("__") and callable(customized_models.__dict__[name]): models.__dict__[name] = customized_models.__dict__[name] model_names = default_model_names + customized_models_names def make_model(args): print("=> creating model '{}'".format(args.arch)) model = models.__dict__[args.arch](progress=True) model.fc = nn.Sequential( nn.Dropout(0.2), nn.Linear(2048, args.num_classes) ) return model if __name__=='__main__': all_model = sorted(name for name in models.__dict__ if not name.startswith("__")) print(all_model) ================================================ FILE: dataset.py ================================================ #!usr/bin/env python #-*- coding:utf-8 _*- """ @version: python3.6 @author: QLMX @contact: wenruichn@gmail.com @time: 2019-09-07 20:27 公众号：AI成长社 知乎：https://www.zhihu.com/people/qlmx-61/columns """ import random import torch from torch.utils.data import Dataset from torch.utils.data import sampler import torchvision.transforms as transforms import pandas as pd import six import sys from PIL import Image import numpy as np class Dataset(Dataset): def __init__(self, root=None, transform=None, target_transform=None, to=None): if '.txt' in root: self.env = list(open(root)) else: self.env = root if not self.env: print('cannot creat lmdb from %s' % (root)) sys.exit(0) self.len = len(self.env) - 1 self.transform = transform self.target_transform = target_transform def __len__(self): return self.len def __getitem__(self, index): assert index <= len(self), 'index range error' index += 1 img_path, label = self.env[index].strip().split(',') try: img = Image.open(img_path) except: print(img_path) print('Corrupted image for %d' % index) return self[index + 1] if self.transform is not None: if img.layers == 1: print(img_path) img = self.transform(img) if self.target_transform is not None: label = self.target_transform(label) return (img, int(label)) class TestDataset(Dataset): def __init__(self, root=None, transform=None, target_transform=None, to=None): if '.txt' in root: self.env = list(open(root)) else: self.env = root if not self.env: print('cannot creat lmdb from %s' % (root)) sys.exit(0) self.len = len(self.env) - 1 self.transform = transform self.target_transform = target_transform self.to = to def __len__(self): return self.len def __getitem__(self, index): assert index <= len(self), 'index range error' index += 1 img_path, label = self.env[index].strip().split(',') try: img = Image.open(img_path) except: print(img_path) print('Corrupted image for %d' % index) return self[index + 1] if self.transform is not None: img = self.transform(img) if self.target_transform is not None: label = self.target_transform(label) return (img, int(label)) class resizeNormalize(object): def __init__(self, size, interpolation=Image.BILINEAR): self.size = size self.interpolation = interpolation self.toTensor = transforms.ToTensor() def __call__(self, img): # padding ratio = self.size[0] / self.size[1] w, h = img.size if w / h < ratio: t = int(h * ratio) w_padding = (t - w) // 2 img = img.crop((-w_padding, 0, w+w_padding, h)) else: t = int(w / ratio) h_padding = (t - h) // 2 img = img.crop((0, -h_padding, w, h+h_padding)) # img.show() # resize img = img.resize(self.size, self.interpolation) img = self.toTensor(img) img.sub_(0.5).div_(0.5) return img ================================================ FILE: models/Res.py ================================================ # Copyright (c) Facebook, Inc. and its affiliates. # All rights reserved. # # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. # import torch.nn as nn try: from torch.hub import load_state_dict_from_url except ImportError: from torch.utils.model_zoo import load_url as load_state_dict_from_url __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152', 'resnext50_32x4d', 'resnext101_32x8d'] model_urls = { 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth', 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', 'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth', 'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth', } def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1): """3x3 convolution with padding""" return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=dilation, groups=groups, bias=False, dilation=dilation) def conv1x1(in_planes, out_planes, stride=1): """1x1 convolution""" return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False) class BasicBlock(nn.Module): expansion = 1 def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, base_width=64, dilation=1, norm_layer=None): super(BasicBlock, self).__init__() if norm_layer is None: norm_layer = nn.BatchNorm2d if groups != 1 or base_width != 64: raise ValueError('BasicBlock only supports groups=1 and base_width=64') if dilation > 1: raise NotImplementedError("Dilation > 1 not supported in BasicBlock") # Both self.conv1 and self.downsample layers downsample the input when stride != 1 self.conv1 = conv3x3(inplanes, planes, stride) self.bn1 = norm_layer(planes) self.relu = nn.ReLU(inplace=True) self.conv2 = conv3x3(planes, planes) self.bn2 = norm_layer(planes) self.downsample = downsample self.stride = stride def forward(self, x): identity = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) if self.downsample is not None: identity = self.downsample(x) out += identity out = self.relu(out) return out class Bottleneck(nn.Module): expansion = 4 def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, base_width=64, dilation=1, norm_layer=None): super(Bottleneck, self).__init__() if norm_layer is None: norm_layer = nn.BatchNorm2d width = int(planes * (base_width / 64.)) * groups # Both self.conv2 and self.downsample layers downsample the input when stride != 1 self.conv1 = conv1x1(inplanes, width) self.bn1 = norm_layer(width) self.conv2 = conv3x3(width, width, stride, groups, dilation) self.bn2 = norm_layer(width) self.conv3 = conv1x1(width, planes * self.expansion) self.bn3 = norm_layer(planes * self.expansion) self.relu = nn.ReLU(inplace=True) self.downsample = downsample self.stride = stride def forward(self, x): identity = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) out = self.relu(out) out = self.conv3(out) out = self.bn3(out) if self.downsample is not None: identity = self.downsample(x) out += identity out = self.relu(out) return out class ResNet(nn.Module): def __init__(self, block, layers, num_classes=1000, zero_init_residual=False, groups=1, width_per_group=64, replace_stride_with_dilation=None, norm_layer=None): super(ResNet, self).__init__() if norm_layer is None: norm_layer = nn.BatchNorm2d self._norm_layer = norm_layer self.inplanes = 64 self.dilation = 1 if replace_stride_with_dilation is None: # each element in the tuple indicates if we should replace # the 2x2 stride with a dilated convolution instead replace_stride_with_dilation = [False, False, False] if len(replace_stride_with_dilation) != 3: raise ValueError("replace_stride_with_dilation should be None " "or a 3-element tuple, got {}".format(replace_stride_with_dilation)) self.groups = groups self.base_width = width_per_group self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False) self.bn1 = norm_layer(self.inplanes) self.relu = nn.ReLU(inplace=True) self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) self.layer1 = self._make_layer(block, 64, layers[0]) self.layer2 = self._make_layer(block, 128, layers[1], stride=2, dilate=replace_stride_with_dilation[0]) self.layer3 = self._make_layer(block, 256, layers[2], stride=2, dilate=replace_stride_with_dilation[1]) self.layer4 = self._make_layer(block, 512, layers[3], stride=2, dilate=replace_stride_with_dilation[2]) self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) self.fc = nn.Linear(512 * block.expansion, num_classes) for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0) # Zero-initialize the last BN in each residual branch, # so that the residual branch starts with zeros, and each residual block behaves like an identity. # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677 if zero_init_residual: for m in self.modules(): if isinstance(m, Bottleneck): nn.init.constant_(m.bn3.weight, 0) elif isinstance(m, BasicBlock): nn.init.constant_(m.bn2.weight, 0) def _make_layer(self, block, planes, blocks, stride=1, dilate=False): norm_layer = self._norm_layer downsample = None previous_dilation = self.dilation if dilate: self.dilation *= stride stride = 1 if stride != 1 or self.inplanes != planes * block.expansion: downsample = nn.Sequential( conv1x1(self.inplanes, planes * block.expansion, stride), norm_layer(planes * block.expansion), ) layers = [] layers.append(block(self.inplanes, planes, stride, downsample, self.groups, self.base_width, previous_dilation, norm_layer)) self.inplanes = planes * block.expansion for _ in range(1, blocks): layers.append(block(self.inplanes, planes, groups=self.groups, base_width=self.base_width, dilation=self.dilation, norm_layer=norm_layer)) return nn.Sequential(*layers) def forward(self, x): x = self.conv1(x) x = self.bn1(x) x = self.relu(x) x = self.maxpool(x) x = self.layer1(x) x = self.layer2(x) x = self.layer3(x) x = self.layer4(x) x = self.avgpool(x) x = x.reshape(x.size(0), -1) x = self.fc(x) return x def _resnet(arch, block, layers, pretrained, progress, **kwargs): model = ResNet(block, layers, **kwargs) if pretrained: state_dict = load_state_dict_from_url(model_urls[arch], progress=progress) model.load_state_dict(state_dict) return model def resnet18(pretrained=False, progress=True, **kwargs): """Constructs a ResNet-18 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress, **kwargs) def resnet34(pretrained=False, progress=True, **kwargs): """Constructs a ResNet-34 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress, **kwargs) def resnet50(pretrained=False, progress=True, **kwargs): """Constructs a ResNet-50 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress, **kwargs) def resnet101(pretrained=False, progress=True, **kwargs): """Constructs a ResNet-101 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs) def resnet152(pretrained=False, progress=True, **kwargs): """Constructs a ResNet-152 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ return _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress, **kwargs) def resnext50_32x4d(pretrained=False, progress=True, **kwargs): """Constructs a ResNeXt-50 32x4d model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ kwargs['groups'] = 32 kwargs['width_per_group'] = 4 return _resnet('resnext50_32x4d', Bottleneck, [3, 4, 6, 3], pretrained, progress, **kwargs) def resnext101_32x8d(pretrained=False, progress=True, **kwargs): """Constructs a ResNeXt-101 32x8d model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ kwargs['groups'] = 32 kwargs['width_per_group'] = 8 return _resnet('resnext101_32x8d', Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs) ================================================ FILE: models/__init__.py ================================================ #!usr/bin/env python #-*- coding:utf-8 _*- """ @version: python3.6 @author: ikkyu-wen @contact: wenruichn@gmail.com @time: 2019-08-17 01:47 """ from __future__ import absolute_import from .resnetxt_wsl import * ================================================ FILE: models/resnet_cbam.py ================================================ #!usr/bin/env python #-*- coding:utf-8 _*- """ @version: python3.6 @author: ikkyu-wen @contact: wenruichn@gmail.com @time: 2019-09-17 12:39 """ import torch.nn as nn import math try: from torch.hub import load_state_dict_from_url except ImportError: from torch.utils.model_zoo import load_url as load_state_dict_from_url import torch __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152', 'resnext50_32x4d', 'resnext101_32x8d'] model_urls = { 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth', 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', 'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth', 'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth', } def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1): """3x3 convolution with padding""" return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=dilation, groups=groups, bias=False, dilation=dilation) def conv1x1(in_planes, out_planes, stride=1): """1x1 convolution""" return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False) class ChannelAttention(nn.Module): def __init__(self, in_planes, ratio=16): super(ChannelAttention, self).__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) self.max_pool = nn.AdaptiveMaxPool2d(1) self.fc1 = nn.Conv2d(in_planes, in_planes // 16, 1, bias=False) self.relu1 = nn.ReLU() self.fc2 = nn.Conv2d(in_planes // 16, in_planes, 1, bias=False) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x)))) max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x)))) out = avg_out + max_out return self.sigmoid(out) class SpatialAttention(nn.Module): def __init__(self, kernel_size=7): super(SpatialAttention, self).__init__() assert kernel_size in (3, 7), 'kernel size must be 3 or 7' padding = 3 if kernel_size == 7 else 1 self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_out = torch.mean(x, dim=1, keepdim=True) max_out, _ = torch.max(x, dim=1, keepdim=True) x = torch.cat([avg_out, max_out], dim=1) x = self.conv1(x) return self.sigmoid(x) class BasicBlock(nn.Module): expansion = 1 def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, base_width=64, dilation=1, norm_layer=None): super(BasicBlock, self).__init__() if norm_layer is None: norm_layer = nn.BatchNorm2d if groups != 1 or base_width != 64: raise ValueError('BasicBlock only supports groups=1 and base_width=64') if dilation > 1: raise NotImplementedError("Dilation > 1 not supported in BasicBlock") # Both self.conv1 and self.downsample layers downsample the input when stride != 1 self.conv1 = conv3x3(inplanes, planes, stride) self.bn1 = norm_layer(planes) self.relu = nn.ReLU(inplace=True) self.conv2 = conv3x3(planes, planes) self.bn2 = norm_layer(planes) self.downsample = downsample self.stride = stride def forward(self, x): identity = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) if self.downsample is not None: identity = self.downsample(x) out += identity out = self.relu(out) return out class Bottleneck(nn.Module): expansion = 4 def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, base_width=64, dilation=1, norm_layer=None): super(Bottleneck, self).__init__() if norm_layer is None: norm_layer = nn.BatchNorm2d width = int(planes * (base_width / 64.)) * groups # Both self.conv2 and self.downsample layers downsample the input when stride != 1 self.conv1 = conv1x1(inplanes, width) self.bn1 = norm_layer(width) self.conv2 = conv3x3(width, width, stride, groups, dilation) self.bn2 = norm_layer(width) self.conv3 = conv1x1(width, planes * self.expansion) self.bn3 = norm_layer(planes * self.expansion) self.relu = nn.ReLU(inplace=True) self.downsample = downsample self.stride = stride def forward(self, x): identity = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) out = self.relu(out) out = self.conv3(out) out = self.bn3(out) if self.downsample is not None: identity = self.downsample(x) out += identity out = self.relu(out) return out class ResNet(nn.Module): def __init__(self, block, layers, num_classes=1000, zero_init_residual=False, groups=1, width_per_group=64, replace_stride_with_dilation=None, norm_layer=None): super(ResNet, self).__init__() if norm_layer is None: norm_layer = nn.BatchNorm2d self._norm_layer = norm_layer self.inplanes = 64 self.dilation = 1 if replace_stride_with_dilation is None: # each element in the tuple indicates if we should replace # the 2x2 stride with a dilated convolution instead replace_stride_with_dilation = [False, False, False] if len(replace_stride_with_dilation) != 3: raise ValueError("replace_stride_with_dilation should be None " "or a 3-element tuple, got {}".format(replace_stride_with_dilation)) self.groups = groups self.base_width = width_per_group self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False) self.bn1 = norm_layer(self.inplanes) self.relu = nn.ReLU(inplace=True) self.ca = ChannelAttention(self.inplanes) self.sa = SpatialAttention() self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) self.layer1 = self._make_layer(block, 64, layers[0]) self.layer2 = self._make_layer(block, 128, layers[1], stride=2, dilate=replace_stride_with_dilation[0]) self.layer3 = self._make_layer(block, 256, layers[2], stride=2, dilate=replace_stride_with_dilation[1]) self.layer4 = self._make_layer(block, 512, layers[3], stride=2, dilate=replace_stride_with_dilation[2]) self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) self.fc = nn.Linear(512 * block.expansion, num_classes) for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0) # Zero-initialize the last BN in each residual branch, # so that the residual branch starts with zeros, and each residual block behaves like an identity. # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677 if zero_init_residual: for m in self.modules(): if isinstance(m, Bottleneck): nn.init.constant_(m.bn3.weight, 0) elif isinstance(m, BasicBlock): nn.init.constant_(m.bn2.weight, 0) def _make_layer(self, block, planes, blocks, stride=1, dilate=False): norm_layer = self._norm_layer downsample = None previous_dilation = self.dilation if dilate: self.dilation *= stride stride = 1 if stride != 1 or self.inplanes != planes * block.expansion: downsample = nn.Sequential( conv1x1(self.inplanes, planes * block.expansion, stride), norm_layer(planes * block.expansion), ) layers = [] layers.append(block(self.inplanes, planes, stride, downsample, self.groups, self.base_width, previous_dilation, norm_layer)) self.inplanes = planes * block.expansion for _ in range(1, blocks): layers.append(block(self.inplanes, planes, groups=self.groups, base_width=self.base_width, dilation=self.dilation, norm_layer=norm_layer)) return nn.Sequential(*layers) def forward(self, x): x = self.conv1(x) x = self.bn1(x) x = self.relu(x) x = self.ca(x) * x x = self.sa(x) * x x = self.maxpool(x) x = self.layer1(x) x = self.layer2(x) x = self.layer3(x) x = self.layer4(x) x = self.avgpool(x) x = x.reshape(x.size(0), -1) x = self.fc(x) return x def _resnet(arch, block, layers, pretrained, progress, **kwargs): model = ResNet(block, layers, **kwargs) if pretrained: state_dict = load_state_dict_from_url(model_urls[arch], progress=progress) new_state_dict = model.state_dict() new_state_dict.update(state_dict) model.load_state_dict(new_state_dict) return model def resnet18(pretrained=False, progress=True, **kwargs): """Constructs a ResNet-18 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress, **kwargs) def resnet34(pretrained=False, progress=True, **kwargs): """Constructs a ResNet-34 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress, **kwargs) def resnet50(pretrained=False, progress=True, **kwargs): """Constructs a ResNet-50 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress, **kwargs) def resnet101(pretrained=False, progress=True, **kwargs): """Constructs a ResNet-101 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs) def resnet152(pretrained=False, progress=True, **kwargs): """Constructs a ResNet-152 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ return _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress, **kwargs) def resnext50_32x4d(pretrained=False, progress=True, **kwargs): """Constructs a ResNeXt-50 32x4d model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ kwargs['groups'] = 32 kwargs['width_per_group'] = 4 return _resnet('resnext50_32x4d', Bottleneck, [3, 4, 6, 3], pretrained, progress, **kwargs) def resnext101_32x8d(pretrained=False, progress=True, **kwargs): """Constructs a ResNeXt-101 32x8d model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ kwargs['groups'] = 32 kwargs['width_per_group'] = 8 return _resnet('resnext101_32x8d', Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs) ================================================ FILE: models/resnetxt_wsl.py ================================================ # Copyright (c) Facebook, Inc. and its affiliates. # All rights reserved. # # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. # Optional list of dependencies required by the package ''' Code From : https://github.com/facebookresearch/WSL-Images/blob/master/hubconf.py ''' __all__ = ['resnext101_32x8d_wsl', 'resnext101_32x16d_wsl', 'resnext101_32x32d_wsl', 'resnext101_32x48d_wsl'] dependencies = ['torch', 'torchvision'] try: from torch.hub import load_state_dict_from_url except ImportError: from torch.utils.model_zoo import load_url as load_state_dict_from_url # from .Res import ResNet, Bottleneck from .resnet_cbam import ResNet, Bottleneck model_urls = { 'resnext101_32x8d': 'https://download.pytorch.org/models/ig_resnext101_32x8-c38310e5.pth', 'resnext101_32x16d': 'https://download.pytorch.org/models/ig_resnext101_32x16-c6f796b0.pth', 'resnext101_32x32d': 'https://download.pytorch.org/models/ig_resnext101_32x32-e4b90b00.pth', 'resnext101_32x48d': 'https://download.pytorch.org/models/ig_resnext101_32x48-3e41cc8a.pth', } # def _resnext(arch, block, layers, pretrained, progress, **kwargs): # model = ResNet(block, layers, **kwargs) # state_dict = load_state_dict_from_url(model_urls[arch], progress=progress) # model.load_state_dict(state_dict) # return model #使用部分加载 def _resnext(arch, block, layers, pretrained, progress, **kwargs): model = ResNet(block, layers, **kwargs) state_dict = load_state_dict_from_url(model_urls[arch], progress=progress) new_state_dict = model.state_dict() new_state_dict.update(state_dict) model.load_state_dict(new_state_dict) return model def resnext101_32x8d_wsl(progress=True, **kwargs): """Constructs a ResNeXt-101 32x8 model pre-trained on weakly-supervised data and finetuned on ImageNet from Figure 5 in `"Exploring the Limits of Weakly Supervised Pretraining" `_ Args: progress (bool): If True, displays a progress bar of the download to stderr. """ kwargs['groups'] = 32 kwargs['width_per_group'] = 8 return _resnext('resnext101_32x8d', Bottleneck, [3, 4, 23, 3], True, progress, **kwargs) def resnext101_32x16d_wsl(progress=True, **kwargs): """Constructs a ResNeXt-101 32x16 model pre-trained on weakly-supervised data and finetuned on ImageNet from Figure 5 in `"Exploring the Limits of Weakly Supervised Pretraining" `_ Args:zz progress (bool): If True, displays a progress bar of the download to stderr. """ kwargs['groups'] = 32 kwargs['width_per_group'] = 16 return _resnext('resnext101_32x16d', Bottleneck, [3, 4, 23, 3], True, progress, **kwargs) def resnext101_32x32d_wsl(progress=True, **kwargs): """Constructs a ResNeXt-101 32x32 model pre-trained on weakly-supervised data and finetuned on ImageNet from Figure 5 in `"Exploring the Limits of Weakly Supervised Pretraining" `_ Args: progress (bool): If True, displays a progress bar of the download to stderr. """ kwargs['groups'] = 32 kwargs['width_per_group'] = 32 return _resnext('resnext101_32x32d', Bottleneck, [3, 4, 23, 3], True, progress, **kwargs) def resnext101_32x48d_wsl(progress=True, **kwargs): """Constructs a ResNeXt-101 32x48 model pre-trained on weakly-supervised data and finetuned on ImageNet from Figure 5 in `"Exploring the Limits of Weakly Supervised Pretraining" `_ Args: progress (bool): If True, displays a progress bar of the download to stderr. """ kwargs['groups'] = 32 kwargs['width_per_group'] = 48 return _resnext('resnext101_32x48d', Bottleneck, [3, 4, 23, 3], True, progress, **kwargs) ================================================ FILE: predict/Res.py ================================================ # Copyright (c) Facebook, Inc. and its affiliates. # All rights reserved. # # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. # import torch.nn as nn try: from torch.hub import load_state_dict_from_url except ImportError: from torch.utils.model_zoo import load_url as load_state_dict_from_url __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152', 'resnext50_32x4d', 'resnext101_32x8d'] model_urls = { 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth', 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', 'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth', 'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth', } def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1): """3x3 convolution with padding""" return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=dilation, groups=groups, bias=False, dilation=dilation) def conv1x1(in_planes, out_planes, stride=1): """1x1 convolution""" return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False) class BasicBlock(nn.Module): expansion = 1 def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, base_width=64, dilation=1, norm_layer=None): super(BasicBlock, self).__init__() if norm_layer is None: norm_layer = nn.BatchNorm2d if groups != 1 or base_width != 64: raise ValueError('BasicBlock only supports groups=1 and base_width=64') if dilation > 1: raise NotImplementedError("Dilation > 1 not supported in BasicBlock") # Both self.conv1 and self.downsample layers downsample the input when stride != 1 self.conv1 = conv3x3(inplanes, planes, stride) self.bn1 = norm_layer(planes) self.relu = nn.ReLU(inplace=True) self.conv2 = conv3x3(planes, planes) self.bn2 = norm_layer(planes) self.downsample = downsample self.stride = stride def forward(self, x): identity = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) if self.downsample is not None: identity = self.downsample(x) out += identity out = self.relu(out) return out class Bottleneck(nn.Module): expansion = 4 def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, base_width=64, dilation=1, norm_layer=None): super(Bottleneck, self).__init__() if norm_layer is None: norm_layer = nn.BatchNorm2d width = int(planes * (base_width / 64.)) * groups # Both self.conv2 and self.downsample layers downsample the input when stride != 1 self.conv1 = conv1x1(inplanes, width) self.bn1 = norm_layer(width) self.conv2 = conv3x3(width, width, stride, groups, dilation) self.bn2 = norm_layer(width) self.conv3 = conv1x1(width, planes * self.expansion) self.bn3 = norm_layer(planes * self.expansion) self.relu = nn.ReLU(inplace=True) self.downsample = downsample self.stride = stride def forward(self, x): identity = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) out = self.relu(out) out = self.conv3(out) out = self.bn3(out) if self.downsample is not None: identity = self.downsample(x) out += identity out = self.relu(out) return out class ResNet(nn.Module): def __init__(self, block, layers, num_classes=1000, zero_init_residual=False, groups=1, width_per_group=64, replace_stride_with_dilation=None, norm_layer=None): super(ResNet, self).__init__() if norm_layer is None: norm_layer = nn.BatchNorm2d self._norm_layer = norm_layer self.inplanes = 64 self.dilation = 1 if replace_stride_with_dilation is None: # each element in the tuple indicates if we should replace # the 2x2 stride with a dilated convolution instead replace_stride_with_dilation = [False, False, False] if len(replace_stride_with_dilation) != 3: raise ValueError("replace_stride_with_dilation should be None " "or a 3-element tuple, got {}".format(replace_stride_with_dilation)) self.groups = groups self.base_width = width_per_group self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False) self.bn1 = norm_layer(self.inplanes) self.relu = nn.ReLU(inplace=True) self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) self.layer1 = self._make_layer(block, 64, layers[0]) self.layer2 = self._make_layer(block, 128, layers[1], stride=2, dilate=replace_stride_with_dilation[0]) self.layer3 = self._make_layer(block, 256, layers[2], stride=2, dilate=replace_stride_with_dilation[1]) self.layer4 = self._make_layer(block, 512, layers[3], stride=2, dilate=replace_stride_with_dilation[2]) self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) self.fc = nn.Linear(512 * block.expansion, num_classes) for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0) # Zero-initialize the last BN in each residual branch, # so that the residual branch starts with zeros, and each residual block behaves like an identity. # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677 if zero_init_residual: for m in self.modules(): if isinstance(m, Bottleneck): nn.init.constant_(m.bn3.weight, 0) elif isinstance(m, BasicBlock): nn.init.constant_(m.bn2.weight, 0) def _make_layer(self, block, planes, blocks, stride=1, dilate=False): norm_layer = self._norm_layer downsample = None previous_dilation = self.dilation if dilate: self.dilation *= stride stride = 1 if stride != 1 or self.inplanes != planes * block.expansion: downsample = nn.Sequential( conv1x1(self.inplanes, planes * block.expansion, stride), norm_layer(planes * block.expansion), ) layers = [] layers.append(block(self.inplanes, planes, stride, downsample, self.groups, self.base_width, previous_dilation, norm_layer)) self.inplanes = planes * block.expansion for _ in range(1, blocks): layers.append(block(self.inplanes, planes, groups=self.groups, base_width=self.base_width, dilation=self.dilation, norm_layer=norm_layer)) return nn.Sequential(*layers) def forward(self, x): x = self.conv1(x) x = self.bn1(x) x = self.relu(x) x = self.maxpool(x) x = self.layer1(x) x = self.layer2(x) x = self.layer3(x) x = self.layer4(x) x = self.avgpool(x) x = x.reshape(x.size(0), -1) x = self.fc(x) return x def _resnet(arch, block, layers, pretrained, progress, **kwargs): model = ResNet(block, layers, **kwargs) if pretrained: state_dict = load_state_dict_from_url(model_urls[arch], progress=progress) model.load_state_dict(state_dict) return model def resnet18(pretrained=False, progress=True, **kwargs): """Constructs a ResNet-18 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress, **kwargs) def resnet34(pretrained=False, progress=True, **kwargs): """Constructs a ResNet-34 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress, **kwargs) def resnet50(pretrained=False, progress=True, **kwargs): """Constructs a ResNet-50 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress, **kwargs) def resnet101(pretrained=False, progress=True, **kwargs): """Constructs a ResNet-101 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs) def resnet152(pretrained=False, progress=True, **kwargs): """Constructs a ResNet-152 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ return _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress, **kwargs) def resnext50_32x4d(pretrained=False, progress=True, **kwargs): """Constructs a ResNeXt-50 32x4d model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ kwargs['groups'] = 32 kwargs['width_per_group'] = 4 return _resnet('resnext50_32x4d', Bottleneck, [3, 4, 6, 3], pretrained, progress, **kwargs) def resnext101_32x8d(pretrained=False, progress=True, **kwargs): """Constructs a ResNeXt-101 32x8d model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ kwargs['groups'] = 32 kwargs['width_per_group'] = 8 return _resnet('resnext101_32x8d', Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs) ================================================ FILE: predict/config.json ================================================ { "model_type": "PyTorch", "runtime":"python3.6", "model_algorithm": "image_classification", "metrics": { "f1": 0.345294, "accuracy": 0.462963, "precision": 0.338977, "recall": 0.351852 }, "apis": [{ "protocol": "http", "url": "/", "method": "post", "request": { "Content-type": "multipart/form-data", "data": { "type": "object", "properties": { "input_img": {"type": "file"} }, "required": ["input_img"] } }, "response": { "Content-type": "multipart/form-data", "data": { "type": "object", "properties": { "result": {"type": "string"} }, "required": ["result"] } } }], "dependencies": [{ "installer": "pip", "packages": [ { "restraint": "ATLEAST", "package_version": "5.2.0", "package_name": "Pillow" } ] }] } ================================================ FILE: predict/customize_service.py ================================================ #!usr/bin/env python #-*- coding:utf-8 _*- """ @version: python3.6 @author: QLMX @contact: wenruichn@gmail.com @time: 2019-08-14 16:07 """ import torchvision.transforms as transforms import torch from torch import nn from PIL import Image import torch.nn.functional as F from resnetxt_wsl import resnext101_32x8d_wsl, resnext101_32x16d_wsl, resnext101_32x32d_wsl from model_service.pytorch_model_service import PTServingBaseService import time from metric.metrics_manager import MetricsManager import log logger = log.getLogger(__name__) args = {} args['arch'] = 'resnext101_32x16d_wsl' args['pretrained'] = False args['num_classes'] = 43 args['big_size'] = 327 args['image_size'] = 288 torch.backends.cudnn.benchmark = True class classfication_service(PTServingBaseService): def __init__(self, model_name, model_path): super(classfication_service, self).__init__(model_name, model_path) self.pre_img = self.preprocess_img1() self.model = self.build_model(model_path) self.model = self.model.cuda() self.model.eval() self.label_id_name_dict = \ { "0": "其他垃圾/一次性快餐盒", "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": "可回收物/纸袋" } def build_model(self, model_path): model = resnext101_32x16d_wsl(pretrained=False, progress=False) model.fc = nn.Sequential( nn.Dropout(0.2), nn.Linear(2048, 43) ) model.load_state_dict(torch.load(model_path)) return model def preprocess_img(self): mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225] infer_transformation = transforms.Compose([ Resize((args['image_size'], args['image_size'])), transforms.ToTensor(), transforms.Normalize(mean=mean, std=std), ]) return infer_transformation def preprocess_img1(self): mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225] return transforms.Compose([ Resize((329, 329)), transforms.CenterCrop(288), transforms.ToTensor(), transforms.Normalize(mean=mean, std=std), ]) def _preprocess(self, data): preprocessed_data = {} for k, v in data.items(): for file_name, file_content in v.items(): img = Image.open(file_content) img = self.pre_img(img) preprocessed_data[k] = img return preprocessed_data def _inference(self, data): img = data['input_img'] img = img.unsqueeze(0).cuda() with torch.no_grad(): pred_score = self.model(img) if pred_score is not None: _, pred_label = torch.max(pred_score.data, 1) return {'result': self.label_id_name_dict[str(pred_label[0].item())]} else: return {'result': 'predict score is None'} # return result def _postprocess(self, data): return data class Resize(object): def __init__(self, size, interpolation=Image.BILINEAR): self.size = size self.interpolation = interpolation def __call__(self, img): ratio = 1 w, h = img.size if w / h < ratio: t = int(h * ratio) w_padding = (t - w) // 2 img = img.crop((-w_padding, 0, w+w_padding, h)) else: t = int(w / ratio) h_padding = (t - h) // 2 img = img.crop((0, -h_padding, w, h+h_padding)) img = img.resize(self.size, self.interpolation) return img ================================================ FILE: predict/predict.py ================================================ #!usr/bin/env python #-*- coding:utf-8 _*- """ @version: python3.6 @author: QLMX @contact: wenruichn@gmail.com @time: 2019-08-14 16:07 """ import torchvision.transforms as transforms import torch from PIL import Image from collections import OrderedDict import torch.nn.functional as F from efficientnet_pytorch import EfficientNet from torch import nn import os, time import torchvision.models as models from resnetxt_wsl import resnext101_32x8d_wsl, resnext101_32x16d_wsl, resnext101_32x32d_wsl os.environ["CUDA_VISIBLE_DEVICES"] = "3" args = {} args['arch'] = 'resnext101_32x16d_wsl' args['pretrained'] = False args['num_classes'] = 43 args['image_size'] = 320 class classfication_service(): def __init__(self, model_path): self.model = self.build_model(model_path) self.pre_img = self.preprocess_img() self.model.eval() self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') self.label_id_name_dict = \ { "0": "其他垃圾/一次性快餐盒", "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": "可回收物/纸袋" } def build_model(self, model_path): if args['arch'] == 'resnext101_32x16d_wsl': model = resnext101_32x16d_wsl(pretrained=False, progress=False) if args['arch'] == 'resnext101_32x8d': model = models.__dict__[args['arch']]() elif args['arch'] == 'efficientnet-b7': model = EfficientNet.from_name(args['arch']) layerName, layer = list(model.named_children())[-1] exec("model." + layerName + "=nn.Linear(layer.in_features," + str(args['num_classes']) + ")") if torch.cuda.is_available(): modelState = torch.load(model_path) model.load_state_dict(modelState) model = model.cuda() else: modelState = torch.load(model_path, map_location='cpu') model.load_state_dict(modelState) return model def preprocess_img(self): mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225] infer_transformation = transforms.Compose([ Resize((args['image_size'], args['image_size'])), transforms.ToTensor(), transforms.Normalize(mean=mean, std=std), ]) return infer_transformation def _preprocess(self, data): preprocessed_data = {} for k, v in data.items(): for file_name, file_content in v.items(): img = Image.open(file_content) img = self.pre_img(img) preprocessed_data[k] = img return preprocessed_data def _inference(self, data): """ model inference function Here are a inference example of resnet, if you use another model, please modify this function """ img = data['input_img'] img = img.unsqueeze(0) img = img.to(self.device) with torch.no_grad(): pred_score = self.model(img) if pred_score is not None: _, pred_label = torch.max(pred_score.data, 1) result = {'result': self.label_id_name_dict[str(pred_label[0].item())]} else: result = {'result': 'predict score is None'} return result def _postprocess(self, data): return data class Resize(object): def __init__(self, size, interpolation=Image.BILINEAR): self.size = size self.interpolation = interpolation def __call__(self, img): ratio = self.size[0] / self.size[1] w, h = img.size if w / h < ratio: t = int(h * ratio) w_padding = (t - w) // 2 img = img.crop((-w_padding, 0, w+w_padding, h)) else: t = int(w / ratio) h_padding = (t - h) // 2 img = img.crop((0, -h_padding, w, h+h_padding)) img = img.resize(self.size, self.interpolation) return img if __name__ == '__main__': model_path = '/Users/QLMX/Documents/' + 'model_19_9992_9590.pth' infer = classfication_service(model_path) input_dir = '/Users/QLMX/Downloads/garbage_classify_v3_select_100' files = os.listdir(input_dir) t1 = int(time.time()*1000) for file_name in files: file_path = os.path.join(input_dir, file_name) img = Image.open(file_path) img = infer.pre_img(img) tt1 = int(time.time() * 1000) result = infer._inference({'input_img': img}) tt2 = int(time.time() * 1000) print((tt2 - tt1) / 100) t2 = int(time.time()*1000) print((t2 - t1)/100) ================================================ FILE: predict/resnet_cbam.py ================================================ #!usr/bin/env python #-*- coding:utf-8 _*- """ @version: python3.6 @author: ikkyu-wen @contact: wenruichn@gmail.com @time: 2019-09-17 12:39 """ import torch.nn as nn import math try: from torch.hub import load_state_dict_from_url except ImportError: from torch.utils.model_zoo import load_url as load_state_dict_from_url import torch __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152', 'resnext50_32x4d', 'resnext101_32x8d'] model_urls = { 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth', 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', 'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth', 'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth', } def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1): """3x3 convolution with padding""" return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=dilation, groups=groups, bias=False, dilation=dilation) def conv1x1(in_planes, out_planes, stride=1): """1x1 convolution""" return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False) class ChannelAttention(nn.Module): def __init__(self, in_planes, ratio=16): super(ChannelAttention, self).__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) self.max_pool = nn.AdaptiveMaxPool2d(1) self.fc1 = nn.Conv2d(in_planes, in_planes // 16, 1, bias=False) self.relu1 = nn.ReLU() self.fc2 = nn.Conv2d(in_planes // 16, in_planes, 1, bias=False) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x)))) max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x)))) out = avg_out + max_out return self.sigmoid(out) class SpatialAttention(nn.Module): def __init__(self, kernel_size=7): super(SpatialAttention, self).__init__() assert kernel_size in (3, 7), 'kernel size must be 3 or 7' padding = 3 if kernel_size == 7 else 1 self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_out = torch.mean(x, dim=1, keepdim=True) max_out, _ = torch.max(x, dim=1, keepdim=True) x = torch.cat([avg_out, max_out], dim=1) x = self.conv1(x) return self.sigmoid(x) class BasicBlock(nn.Module): expansion = 1 def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, base_width=64, dilation=1, norm_layer=None): super(BasicBlock, self).__init__() if norm_layer is None: norm_layer = nn.BatchNorm2d if groups != 1 or base_width != 64: raise ValueError('BasicBlock only supports groups=1 and base_width=64') if dilation > 1: raise NotImplementedError("Dilation > 1 not supported in BasicBlock") # Both self.conv1 and self.downsample layers downsample the input when stride != 1 self.conv1 = conv3x3(inplanes, planes, stride) self.bn1 = norm_layer(planes) self.relu = nn.ReLU(inplace=True) self.conv2 = conv3x3(planes, planes) self.bn2 = norm_layer(planes) self.downsample = downsample self.stride = stride def forward(self, x): identity = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) if self.downsample is not None: identity = self.downsample(x) out += identity out = self.relu(out) return out class Bottleneck(nn.Module): expansion = 4 def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, base_width=64, dilation=1, norm_layer=None): super(Bottleneck, self).__init__() if norm_layer is None: norm_layer = nn.BatchNorm2d width = int(planes * (base_width / 64.)) * groups # Both self.conv2 and self.downsample layers downsample the input when stride != 1 self.conv1 = conv1x1(inplanes, width) self.bn1 = norm_layer(width) self.conv2 = conv3x3(width, width, stride, groups, dilation) self.bn2 = norm_layer(width) self.conv3 = conv1x1(width, planes * self.expansion) self.bn3 = norm_layer(planes * self.expansion) self.relu = nn.ReLU(inplace=True) self.downsample = downsample self.stride = stride def forward(self, x): identity = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) out = self.relu(out) out = self.conv3(out) out = self.bn3(out) if self.downsample is not None: identity = self.downsample(x) out += identity out = self.relu(out) return out class ResNet(nn.Module): def __init__(self, block, layers, num_classes=1000, zero_init_residual=False, groups=1, width_per_group=64, replace_stride_with_dilation=None, norm_layer=None): super(ResNet, self).__init__() if norm_layer is None: norm_layer = nn.BatchNorm2d self._norm_layer = norm_layer self.inplanes = 64 self.dilation = 1 if replace_stride_with_dilation is None: # each element in the tuple indicates if we should replace # the 2x2 stride with a dilated convolution instead replace_stride_with_dilation = [False, False, False] if len(replace_stride_with_dilation) != 3: raise ValueError("replace_stride_with_dilation should be None " "or a 3-element tuple, got {}".format(replace_stride_with_dilation)) self.groups = groups self.base_width = width_per_group self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False) self.bn1 = norm_layer(self.inplanes) self.relu = nn.ReLU(inplace=True) self.ca = ChannelAttention(self.inplanes) self.sa = SpatialAttention() self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) self.layer1 = self._make_layer(block, 64, layers[0]) self.layer2 = self._make_layer(block, 128, layers[1], stride=2, dilate=replace_stride_with_dilation[0]) self.layer3 = self._make_layer(block, 256, layers[2], stride=2, dilate=replace_stride_with_dilation[1]) self.layer4 = self._make_layer(block, 512, layers[3], stride=2, dilate=replace_stride_with_dilation[2]) self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) self.fc = nn.Linear(512 * block.expansion, num_classes) for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0) # Zero-initialize the last BN in each residual branch, # so that the residual branch starts with zeros, and each residual block behaves like an identity. # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677 if zero_init_residual: for m in self.modules(): if isinstance(m, Bottleneck): nn.init.constant_(m.bn3.weight, 0) elif isinstance(m, BasicBlock): nn.init.constant_(m.bn2.weight, 0) def _make_layer(self, block, planes, blocks, stride=1, dilate=False): norm_layer = self._norm_layer downsample = None previous_dilation = self.dilation if dilate: self.dilation *= stride stride = 1 if stride != 1 or self.inplanes != planes * block.expansion: downsample = nn.Sequential( conv1x1(self.inplanes, planes * block.expansion, stride), norm_layer(planes * block.expansion), ) layers = [] layers.append(block(self.inplanes, planes, stride, downsample, self.groups, self.base_width, previous_dilation, norm_layer)) self.inplanes = planes * block.expansion for _ in range(1, blocks): layers.append(block(self.inplanes, planes, groups=self.groups, base_width=self.base_width, dilation=self.dilation, norm_layer=norm_layer)) return nn.Sequential(*layers) def forward(self, x): x = self.conv1(x) x = self.bn1(x) x = self.relu(x) x = self.ca(x) * x x = self.sa(x) * x x = self.maxpool(x) x = self.layer1(x) x = self.layer2(x) x = self.layer3(x) x = self.layer4(x) x = self.avgpool(x) x = x.reshape(x.size(0), -1) x = self.fc(x) return x def _resnet(arch, block, layers, pretrained, progress, **kwargs): model = ResNet(block, layers, **kwargs) if pretrained: state_dict = load_state_dict_from_url(model_urls[arch], progress=progress) new_state_dict = model.state_dict() new_state_dict.update(state_dict) model.load_state_dict(new_state_dict) return model def resnet18(pretrained=False, progress=True, **kwargs): """Constructs a ResNet-18 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress, **kwargs) def resnet34(pretrained=False, progress=True, **kwargs): """Constructs a ResNet-34 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress, **kwargs) def resnet50(pretrained=False, progress=True, **kwargs): """Constructs a ResNet-50 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress, **kwargs) def resnet101(pretrained=False, progress=True, **kwargs): """Constructs a ResNet-101 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs) def resnet152(pretrained=False, progress=True, **kwargs): """Constructs a ResNet-152 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ return _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress, **kwargs) def resnext50_32x4d(pretrained=False, progress=True, **kwargs): """Constructs a ResNeXt-50 32x4d model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ kwargs['groups'] = 32 kwargs['width_per_group'] = 4 return _resnet('resnext50_32x4d', Bottleneck, [3, 4, 6, 3], pretrained, progress, **kwargs) def resnext101_32x8d(pretrained=False, progress=True, **kwargs): """Constructs a ResNeXt-101 32x8d model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """ kwargs['groups'] = 32 kwargs['width_per_group'] = 8 return _resnet('resnext101_32x8d', Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs) ================================================ FILE: predict/resnetxt_wsl.py ================================================ # Copyright (c) Facebook, Inc. and its affiliates. # All rights reserved. # # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. # Optional list of dependencies required by the package ''' Code From : https://github.com/facebookresearch/WSL-Images/blob/master/hubconf.py ''' __all__ = ['resnext101_32x8d_wsl', 'resnext101_32x16d_wsl', 'resnext101_32x32d_wsl', 'resnext101_32x48d_wsl'] dependencies = ['torch', 'torchvision'] try: from torch.hub import load_state_dict_from_url except ImportError: from torch.utils.model_zoo import load_url as load_state_dict_from_url # from Res import ResNet, Bottleneck from .resnet_cbam import ResNet, Bottleneck model_urls = { 'resnext101_32x8d': 'https://download.pytorch.org/models/ig_resnext101_32x8-c38310e5.pth', 'resnext101_32x16d': 'https://download.pytorch.org/models/ig_resnext101_32x16-c6f796b0.pth', 'resnext101_32x32d': 'https://download.pytorch.org/models/ig_resnext101_32x32-e4b90b00.pth', 'resnext101_32x48d': 'https://download.pytorch.org/models/ig_resnext101_32x48-3e41cc8a.pth', } # def _resnext(arch, block, layers, pretrained, progress, **kwargs): # model = ResNet(block, layers, **kwargs) # if pretrained: # state_dict = load_state_dict_from_url(model_urls[arch], progress=progress) # model.load_state_dict(state_dict) # return model #使用部分加载 def _resnext(arch, block, layers, pretrained, progress, **kwargs): model = ResNet(block, layers, **kwargs) if pretrained: state_dict = load_state_dict_from_url(model_urls[arch], progress=progress) new_state_dict = model.state_dict() new_state_dict.update(state_dict) model.load_state_dict(new_state_dict) return model def resnext101_32x8d_wsl(pretrained=True, progress=True, **kwargs): """Constructs a ResNeXt-101 32x8 model pre-trained on weakly-supervised data and finetuned on ImageNet from Figure 5 in `"Exploring the Limits of Weakly Supervised Pretraining" `_ Args: progress (bool): If True, displays a progress bar of the download to stderr. """ kwargs['groups'] = 32 kwargs['width_per_group'] = 8 return _resnext('resnext101_32x8d', Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs) def resnext101_32x16d_wsl(pretrained=True, progress=True, **kwargs): """Constructs a ResNeXt-101 32x16 model pre-trained on weakly-supervised data and finetuned on ImageNet from Figure 5 in `"Exploring the Limits of Weakly Supervised Pretraining" `_ Args: progress (bool): If True, displays a progress bar of the download to stderr. """ kwargs['groups'] = 32 kwargs['width_per_group'] = 16 return _resnext('resnext101_32x16d', Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs) def resnext101_32x32d_wsl(pretrained=True, progress=True, **kwargs): """Constructs a ResNeXt-101 32x32 model pre-trained on weakly-supervised data and finetuned on ImageNet from Figure 5 in `"Exploring the Limits of Weakly Supervised Pretraining" `_ Args: progress (bool): If True, displays a progress bar of the download to stderr. """ kwargs['groups'] = 32 kwargs['width_per_group'] = 32 return _resnext('resnext101_32x32d', Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs) def resnext101_32x48d_wsl(pretrained=True, progress=True, **kwargs): """Constructs a ResNeXt-101 32x48 model pre-trained on weakly-supervised data and finetuned on ImageNet from Figure 5 in `"Exploring the Limits of Weakly Supervised Pretraining" `_ Args: progress (bool): If True, displays a progress bar of the download to stderr. """ kwargs['groups'] = 32 kwargs['width_per_group'] = 48 return _resnext('resnext101_32x48d', Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs) ================================================ FILE: preprocess.py ================================================ #!usr/bin/env python #-*- coding:utf-8 _*- """ @version: python3.6 @author: QLMX @contact: wenruichn@gmail.com @time: 2019-08-14 11:16 公众号：AI成长社 知乎：https://www.zhihu.com/people/qlmx-61/columns """ from glob import glob import os import codecs import random import numpy as np from sklearn.model_selection import KFold, StratifiedKFold base_path = 'data/' data_path = base_path + 'garbage_classify/train_data' label_files = glob(os.path.join(data_path, '*.txt')) img_paths = [] labels = [] result = [] label_dict = {} data_dict = {} for index, file_path in enumerate(label_files): with codecs.open(file_path, 'r', 'utf-8') as f: line = f.readline() line_split = line.strip().split(', ') if len(line_split) != 2: print('%s contain error lable' % os.path.basename(file_path)) continue img_name = line_split[0] label = int(line_split[1]) img_paths.append(os.path.join(data_path, img_name)) labels.append(label) result.append(os.path.join(data_path, img_name) + ',' + str(label)) label_dict[label] = label_dict.get(label, 0) + 1 if label not in data_dict: data_dict[label] = [] data_dict[label].append(os.path.join(data_path, img_name) + ',' + str(label)) data_path_add = base_path + 'garbage_classify_v3' label_files_add = glob(os.path.join(data_path_add, '*.txt')) for index, file_path in enumerate(label_files_add): with codecs.open(file_path, 'r', 'utf-8') as f: line = f.readline() line_split = line.strip().split(', ') if len(line_split) != 2: print('%s contain error lable' % os.path.basename(file_path)) continue img_name = line_split[0] label = int(line_split[1]) img_paths.append(os.path.join(data_path_add, img_name)) labels.append(label) result.append(os.path.join(data_path_add, img_name) + ',' + str(label)) label_dict[label] = label_dict.get(label, 0) + 1 if label not in data_dict: data_dict[label] = [] data_dict[label].append(os.path.join(data_path_add, img_name) + ',' + str(label)) folds = StratifiedKFold(n_splits=10, shuffle=True, random_state=2019) for fold_, (trn_idx, val_idx) in enumerate(folds.split(result, labels)): train_data = list(np.array(result)[trn_idx]) val_data = list(np.array(result)[val_idx]) print(len(train_data), len(val_data)) with open(base_path + 'train1.txt', 'w') as f1: for item in train_data: f1.write(item + '\n') with open(base_path + 'val1.txt', 'w') as f2: for item in val_data: f2.write(item + '\n') from PIL import Image ###predata 2 all_data = [] train = [] val = [] rate = 0.9 import cv2 from tqdm import tqdm error_list = ['data/additional_train_data/38/242.jpg', 'data/additional_train_data/34/79.jpg', 'data/additional_train_data/27/55.jpg' 'data/new/8/0.jpg' ] data_path = base_path + 'new/' for i in range(40): na_item = [] img_files = glob(os.path.join(data_path, str(i), '*.jpg')) for item in tqdm(img_files): ii = cv2.imread(item) if item not in error_list: jj = Image.open(item).layers if jj == 1: print(item) all_data.append(item + ',' + str(i)) na_item.append(item + ',' + str(i)) random.shuffle(na_item) train.extend(na_item[ : int(len(na_item)*rate)]) val.extend(na_item[int(len(na_item)*rate):]) print(len(train), len(val)) random.shuffle(all_data) random.shuffle(train) random.shuffle(val) print(len(all_data)) old = [] with open(base_path + 'train1.txt', 'r') as f: for i in f.readlines(): old.append(i.strip()) for i in all_data: img_path, label = i.strip().split(',') all_data.extend(old) print(len(all_data)) random.shuffle(all_data) with open(base_path + 'new_shu_label.txt', 'w') as f1: for item in all_data: f1.write(item + '\n') ================================================ FILE: requirements.txt ================================================ pytorch==1.0.1 torchvision==0.2.2 matplotlib==3.1.0 numpy==1.16.4 scikit-image pandas sklearn ================================================ FILE: train.py ================================================ #!usr/bin/env python #-*- coding:utf-8 _*- """ @version: python3.6 @author: QLMX @contact: wenruichn@gmail.com @time: 2019-08-14 11:08 公众号：AI成长社 知乎：https://www.zhihu.com/people/qlmx-61/columns """ from __future__ import print_function import os import time import random import torch import torch.nn as nn import torch.nn.parallel import torch.backends.cudnn as cudnn import torch.utils.data as data import torchvision.transforms as transforms import dataset import numpy as np from args import args from build_net import make_model from transform import get_transforms from utils import Bar, Logger, AverageMeter, accuracy, mkdir_p, savefig, get_optimizer, save_checkpoint state = {k: v for k, v in args._get_kwargs()} # Use CUDA os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id use_cuda = torch.cuda.is_available() # Random seed if args.manualSeed is None: args.manualSeed = random.randint(1, 10000) random.seed(args.manualSeed) torch.manual_seed(args.manualSeed) if use_cuda: torch.cuda.manual_seed_all(args.manualSeed) best_acc = 0 # best test accuracy def main(): global best_acc start_epoch = args.start_epoch # start from epoch 0 or last checkpoint epoch if not os.path.isdir(args.checkpoint): mkdir_p(args.checkpoint) # Data transform = get_transforms(input_size=args.image_size, test_size=args.image_size, backbone=None) print('==> Preparing dataset %s' % args.trainroot) trainset = dataset.Dataset(root=args.trainroot, transform=transform['val_train']) train_loader = data.DataLoader(trainset, batch_size=args.train_batch, shuffle=True, num_workers=args.workers, pin_memory=True) valset = dataset.TestDataset(root=args.valroot, transform=transform['val_test']) val_loader = data.DataLoader(valset, batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=True) model = make_model(args) if args.arch.startswith('alexnet') or args.arch.startswith('vgg'): model.features = torch.nn.DataParallel(model.features) model.cuda() else: model = torch.nn.DataParallel(model).cuda() cudnn.benchmark = True print(' Total params: %.2fM' % (sum(p.numel() for p in model.parameters())/1000000.0)) # define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss().cuda() optimizer = get_optimizer(model, args) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, patience=5, verbose=False) # Resume title = 'ImageNet-' + args.arch if args.resume: # Load checkpoint. print('==> Resuming from checkpoint..') assert os.path.isfile(args.resume), 'Error: no checkpoint directory found!' args.checkpoint = os.path.dirname(args.resume) checkpoint = torch.load(args.resume) best_acc = checkpoint['best_acc'] start_epoch = checkpoint['epoch'] model.module.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title, resume=True) else: logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title) logger.set_names(['Learning Rate', 'Train Loss', 'Valid Loss', 'Train Acc.', 'Valid Acc.']) if args.evaluate: print('\nEvaluation only') test_loss, test_acc = test(val_loader, model, criterion, start_epoch, use_cuda) print(' Test Loss: %.8f, Test Acc: %.2f' % (test_loss, test_acc)) return # Train and val for epoch in range(start_epoch, args.epochs): print('\nEpoch: [%d | %d] LR: %f' % (epoch + 1, args.epochs, optimizer.param_groups[0]['lr'])) train_loss, train_acc, train_5 = train(train_loader, model, criterion, optimizer, epoch, use_cuda) test_loss, test_acc, test_5 = test(val_loader, model, criterion, epoch, use_cuda) scheduler.step(test_loss) # append logger file logger.append([state['lr'], train_loss, test_loss, train_acc, test_acc]) print('train_loss:%f, val_loss:%f, train_acc:%f, train_5:%f, val_acc:%f, val_5:%f' % (train_loss, test_loss, train_acc, train_5, test_acc, test_5)) # save model is_best = test_acc > best_acc best_acc = max(test_acc, best_acc) if len(args.gpu_id) > 1: save_checkpoint({ 'fold': 0, 'epoch': epoch + 1, 'state_dict': model.module.state_dict(), 'train_acc': train_acc, 'acc': test_acc, 'best_acc': best_acc, 'optimizer': optimizer.state_dict(), }, is_best, single=True, checkpoint=args.checkpoint) else: save_checkpoint({ 'fold': 0, 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'train_acc':train_acc, 'acc': test_acc, 'best_acc': best_acc, 'optimizer' : optimizer.state_dict(), }, is_best, single=True, checkpoint=args.checkpoint) logger.close() logger.plot() savefig(os.path.join(args.checkpoint, 'log.eps')) print('Best acc:') print(best_acc) def train(train_loader, model, criterion, optimizer, epoch, use_cuda): # switch to train mode model.train() batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() top1 = AverageMeter() top5 = AverageMeter() end = time.time() bar = Bar('Processing', max=len(train_loader)) for batch_idx, (inputs, targets) in enumerate(train_loader): # measure data loading time data_time.update(time.time() - end) if use_cuda: inputs, targets = inputs.cuda(), targets.cuda(async=True) inputs, targets = torch.autograd.Variable(inputs), torch.autograd.Variable(targets) # compute output outputs = model(inputs) loss = criterion(outputs, targets) # measure accuracy and record loss prec1, prec5 = accuracy(outputs.data, targets.data, topk=(1, 5)) losses.update(loss.item(), inputs.size(0)) top1.update(prec1.item(), inputs.size(0)) top5.update(prec5.item(), inputs.size(0)) # compute gradient and do SGD step optimizer.zero_grad() loss.backward() optimizer.step() # measure elapsed time batch_time.update(time.time() - end) end = time.time() # plot progress bar.suffix = '({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | top1: {top1: .4f} | top5: {top5: .4f}'.format( batch=batch_idx + 1, size=len(train_loader), data=data_time.val, bt=batch_time.val, total=bar.elapsed_td, eta=bar.eta_td, loss=losses.avg, top1=top1.avg, top5=top5.avg, ) bar.next() bar.finish() return (losses.avg, top1.avg, top5.avg) def test(val_loader, model, criterion, epoch, use_cuda): global best_acc batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() top1 = AverageMeter() top5 = AverageMeter() # switch to evaluate mode model.eval() end = time.time() bar = Bar('Processing', max=len(val_loader)) for batch_idx, (inputs, targets) in enumerate(val_loader): # measure data loading time data_time.update(time.time() - end) if use_cuda: inputs, targets = inputs.cuda(), targets.cuda() inputs, targets = torch.autograd.Variable(inputs), torch.autograd.Variable(targets) # compute output outputs = model(inputs) loss = criterion(outputs, targets) # measure accuracy and record loss prec1, prec5 = accuracy(outputs.data, targets.data, topk=(1, 5)) losses.update(loss.item(), inputs.size(0)) top1.update(prec1.item(), inputs.size(0)) top5.update(prec5.item(), inputs.size(0)) # measure elapsed time batch_time.update(time.time() - end) end = time.time() # plot progress bar.suffix = '({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | top1: {top1: .4f} | top5: {top5: .4f}'.format( batch=batch_idx + 1, size=len(val_loader), data=data_time.avg, bt=batch_time.avg, total=bar.elapsed_td, eta=bar.eta_td, loss=losses.avg, top1=top1.avg, top5=top5.avg, ) bar.next() bar.finish() return (losses.avg, top1.avg, top5.avg) if __name__ == '__main__': main() ================================================ FILE: transform.py ================================================ #!usr/bin/env python #-*- coding:utf-8 _*- """ @version: python3.6 @author: QLMX @contact: wenruichn@gmail.com @time: 2019-09-07 18:54 公众号：AI成长社 知乎：https://www.zhihu.com/people/qlmx-61/columns """ import random import math import torch from PIL import Image, ImageOps, ImageFilter from torchvision import transforms class Resize(object): def __init__(self, size, interpolation=Image.BILINEAR): self.size = size self.interpolation = interpolation def __call__(self, img): # padding ratio = self.size[0] / self.size[1] w, h = img.size if w / h < ratio: t = int(h * ratio) w_padding = (t - w) // 2 img = img.crop((-w_padding, 0, w+w_padding, h)) else: t = int(w / ratio) h_padding = (t - h) // 2 img = img.crop((0, -h_padding, w, h+h_padding)) img = img.resize(self.size, self.interpolation) return img class RandomRotate(object): def __init__(self, degree, p=0.5): self.degree = degree self.p = p def __call__(self, img): if random.random() < self.p: rotate_degree = random.uniform(-1*self.degree, self.degree) img = img.rotate(rotate_degree, Image.BILINEAR) return img class RandomGaussianBlur(object): def __init__(self, p=0.5): self.p = p def __call__(self, img): if random.random() < self.p: img = img.filter(ImageFilter.GaussianBlur( radius=random.random())) return img def get_train_transform(mean, std, size): train_transform = transforms.Compose([ Resize((int(size * (256 / 224)), int(size * (256 / 224)))), transforms.RandomCrop(size), transforms.RandomHorizontalFlip(), # RandomRotate(15, 0.3), # RandomGaussianBlur(), transforms.ToTensor(), transforms.Normalize(mean=mean, std=std), ]) return train_transform def get_test_transform(mean, std, size): return transforms.Compose([ Resize((int(size * (256 / 224)), int(size * (256 / 224)))), transforms.CenterCrop(size), transforms.ToTensor(), transforms.Normalize(mean=mean, std=std), ]) def get_transforms(input_size=224, test_size=224, backbone=None): mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225] if backbone is not None and backbone in ['pnasnet5large', 'nasnetamobile']: mean, std = [0.5, 0.5, 0.5], [0.5, 0.5, 0.5] transformations = {} transformations['val_train'] = get_train_transform(mean, std, input_size) transformations['val_test'] = get_test_transform(mean, std, test_size) return transformations ================================================ FILE: utils/__init__.py ================================================ """Useful utils """ from .misc import * from .logger import * from .visualize import * from .eval import * from .utils import * from .radam import * # progress bar import os, sys sys.path.append(os.path.join(os.path.dirname(__file__), "progress")) print(os.path.dirname(__file__)) from progress.bar import Bar as Bar ================================================ FILE: utils/eval.py ================================================ from __future__ import print_function, absolute_import __all__ = ['accuracy', 'precision'] def accuracy(output, target, topk=(1,)): """Computes the precision@k for the specified values of k""" maxk = max(topk) batch_size = target.size(0) _, pred = output.topk(maxk, 1, True, True) pred = pred.t() correct = pred.eq(target.view(1, -1).expand_as(pred)) res = [] for k in topk: correct_k = correct[:k].view(-1).float().sum(0) res.append(correct_k.mul_(100.0 / batch_size)) return res def precision(output, target): pass ================================================ FILE: utils/logger.py ================================================ # A simple torch style logger # (C) Wei YANG 2017 from __future__ import absolute_import import matplotlib.pyplot as plt import os import sys import numpy as np __all__ = ['Logger', 'LoggerMonitor', 'savefig'] def savefig(fname, dpi=None): dpi = 150 if dpi == None else dpi plt.savefig(fname, dpi=dpi) def plot_overlap(logger, names=None): names = logger.names if names == None else names numbers = logger.numbers for _, name in enumerate(names): x = np.arange(len(numbers[name])) plt.plot(x, np.asarray(numbers[name])) return [logger.title + '(' + name + ')' for name in names] class Logger(object): '''Save training process to log file with simple plot function.''' def __init__(self, fpath, title=None, resume=False): self.file = None self.resume = resume self.title = '' if title == None else title if fpath is not None: if resume: self.file = open(fpath, 'r') name = self.file.readline() self.names = name.rstrip().split('\t') self.numbers = {} for _, name in enumerate(self.names): self.numbers[name] = [] for numbers in self.file: numbers = numbers.rstrip().split('\t') for i in range(0, len(numbers)): self.numbers[self.names[i]].append(numbers[i]) self.file.close() self.file = open(fpath, 'a') else: self.file = open(fpath, 'w') def set_names(self, names): if self.resume: pass # initialize numbers as empty list self.numbers = {} self.names = names for _, name in enumerate(self.names): self.file.write(name) self.file.write('\t') self.numbers[name] = [] self.file.write('\n') self.file.flush() def append(self, numbers): assert len(self.names) == len(numbers), 'Numbers do not match names' for index, num in enumerate(numbers): self.file.write("{0:.6f}".format(num)) self.file.write('\t') self.numbers[self.names[index]].append(num) self.file.write('\n') self.file.flush() def plot(self, names=None): names = self.names if names == None else names numbers = self.numbers for _, name in enumerate(names): x = np.arange(len(numbers[name])) plt.plot(x, np.asarray(numbers[name])) plt.legend([self.title + '(' + name + ')' for name in names]) plt.grid(True) def close(self): if self.file is not None: self.file.close() class LoggerMonitor(object): '''Load and visualize multiple logs.''' def __init__ (self, paths): '''paths is a distionary with {name:filepath} pair''' self.loggers = [] for title, path in paths.items(): logger = Logger(path, title=title, resume=True) self.loggers.append(logger) def plot(self, names=None): plt.figure() plt.subplot(121) legend_text = [] for logger in self.loggers: legend_text += plot_overlap(logger, names) plt.legend(legend_text, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.) plt.grid(True) if __name__ == '__main__': # # Example # logger = Logger('test.txt') # logger.set_names(['Train loss', 'Valid loss','Test loss']) # length = 100 # t = np.arange(length) # train_loss = np.exp(-t / 10.0) + np.random.rand(length) * 0.1 # valid_loss = np.exp(-t / 10.0) + np.random.rand(length) * 0.1 # test_loss = np.exp(-t / 10.0) + np.random.rand(length) * 0.1 # for i in range(0, length): # logger.append([train_loss[i], valid_loss[i], test_loss[i]]) # logger.plot() # Example: logger monitor paths = { 'resadvnet20':'/home/wyang/code/pytorch-classification/checkpoint/cifar10/resadvnet20/log.txt', 'resadvnet32':'/home/wyang/code/pytorch-classification/checkpoint/cifar10/resadvnet32/log.txt', 'resadvnet44':'/home/wyang/code/pytorch-classification/checkpoint/cifar10/resadvnet44/log.txt', } field = ['Valid Acc.'] monitor = LoggerMonitor(paths) monitor.plot(names=field) savefig('test.eps') ================================================ FILE: utils/misc.py ================================================ '''Some helper functions for PyTorch, including: - get_mean_and_std: calculate the mean and std value of dataset. - msr_init: net parameter initialization. - progress_bar: progress bar mimic xlua.progress. ''' import errno import os import sys import time import math import torch.nn as nn import torch.nn.init as init from torch.autograd import Variable import torch import shutil import adabound from utils.radam import RAdam, AdamW import torchvision.transforms as transforms __all__ = ['get_mean_and_std', 'init_params', 'mkdir_p', 'AverageMeter', 'get_optimizer', 'save_checkpoint'] def get_mean_and_std(dataset): '''Compute the mean and std value of dataset.''' dataloader = trainloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=2) mean = torch.zeros(3) std = torch.zeros(3) print('==> Computing mean and std..') for inputs, targets in dataloader: for i in range(3): mean[i] += inputs[:,i,:,:].mean() std[i] += inputs[:,i,:,:].std() mean.div_(len(dataset)) std.div_(len(dataset)) return mean, std def init_params(net): '''Init layer parameters.''' for m in net.modules(): if isinstance(m, nn.Conv2d): init.kaiming_normal(m.weight, mode='fan_out') if m.bias: init.constant(m.bias, 0) elif isinstance(m, nn.BatchNorm2d): init.constant(m.weight, 1) init.constant(m.bias, 0) elif isinstance(m, nn.Linear): init.normal(m.weight, std=1e-3) if m.bias: init.constant(m.bias, 0) def mkdir_p(path): '''make dir if not exist''' try: os.makedirs(path) except OSError as exc: # Python >2.5 if exc.errno == errno.EEXIST and os.path.isdir(path): pass else: raise class AverageMeter(object): """Computes and stores the average and current value Imported from https://github.com/pytorch/examples/blob/master/imagenet/main.py#L247-L262 """ def __init__(self): self.reset() def reset(self): self.val = 0 self.avg = 0 self.sum = 0 self.count = 0 def update(self, val, n=1): self.val = val self.sum += val * n self.count += n self.avg = self.sum / self.count def get_optimizer(model, args): parameters = [] for name, param in model.named_parameters(): if 'fc' in name or 'class' in name or 'last_linear' in name or 'ca' in name or 'sa' in name: parameters.append({'params': param, 'lr': args.lr * args.lr_fc_times}) else: parameters.append({'params': param, 'lr': args.lr}) if args.optimizer == 'sgd': return torch.optim.SGD(parameters, # model.parameters(), args.lr, momentum=args.momentum, nesterov=args.nesterov, weight_decay=args.weight_decay) elif args.optimizer == 'rmsprop': return torch.optim.RMSprop(parameters, # model.parameters(), args.lr, alpha=args.alpha, weight_decay=args.weight_decay) elif args.optimizer == 'adam': return torch.optim.Adam(parameters, # model.parameters(), args.lr, betas=(args.beta1, args.beta2), weight_decay=args.weight_decay) elif args.optimizer == 'AdaBound': return adabound.AdaBound(parameters, # model.parameters(), lr=args.lr, final_lr=args.final_lr) elif args.optimizer == 'radam': return RAdam(parameters, lr=args.lr, betas=(args.beta1, args.beta2), weight_decay=args.weight_decay) else: raise NotImplementedError def save_checkpoint(state, is_best, single=True, checkpoint='checkpoint', filename='checkpoint.pth.tar'): if single: fold = '' else: fold = str(state['fold']) + '_' cur_name = 'checkpoint.pth.tar' filepath = os.path.join(checkpoint, fold + cur_name) curpath = os.path.join(checkpoint, fold + 'model_cur.pth') torch.save(state, filepath) torch.save(state['state_dict'], curpath) if is_best and state['epoch'] >= 5: model_name = 'model_' + str(state['epoch']) + '_' + str(int(round(state['train_acc']*100, 0))) + '_' + str(int(round(state['acc']*100, 0))) + '.pth' model_path = os.path.join(checkpoint, fold + model_name) torch.save(state['state_dict'], model_path) def save_checkpoint2(state, is_best, checkpoint='checkpoint', filename='checkpoint.pth.tar'): # best_model = '/application/search/qlmx/clover/garbage/code/image_classfication/predict/' fold = str(state['fold']) + '_' filepath = os.path.join(checkpoint, fold + filename) model_path = os.path.join(checkpoint, fold + 'model_cur.pth') torch.save(state, filepath) torch.save(state['state_dict'], model_path) if is_best: shutil.copyfile(filepath, os.path.join(checkpoint, fold + 'model_best.pth.tar')) shutil.copyfile(model_path, os.path.join(checkpoint, fold + 'model_best.pth')) ================================================ FILE: utils/progress/LICENSE ================================================ # Copyright (c) 2012 Giorgos Verigakis # # Permission to use, copy, modify, and distribute this software for any # purpose with or without fee is hereby granted, provided that the above # copyright notice and this permission notice appear in all copies. # # THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES # WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF # MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR # ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES # WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN # ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF # OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. ================================================ FILE: utils/progress/MANIFEST.in ================================================ include README.rst LICENSE ================================================ FILE: utils/progress/README.rst ================================================ Easy progress reporting for Python ================================== |pypi| |demo| .. |pypi| image:: https://img.shields.io/pypi/v/progress.svg .. |demo| image:: https://raw.github.com/verigak/progress/master/demo.gif :alt: Demo Bars ---- There are 7 progress bars to choose from: - ``Bar`` - ``ChargingBar`` - ``FillingSquaresBar`` - ``FillingCirclesBar`` - ``IncrementalBar`` - ``PixelBar`` - ``ShadyBar`` To use them, just call ``next`` to advance and ``finish`` to finish: .. code-block:: python from progress.bar import Bar bar = Bar('Processing', max=20) for i in range(20): # Do some work bar.next() bar.finish() The result will be a bar like the following: :: Processing |############# | 42/100 To simplify the common case where the work is done in an iterator, you can use the ``iter`` method: .. code-block:: python for i in Bar('Processing').iter(it): # Do some work Progress bars are very customizable, you can change their width, their fill character, their suffix and more: .. code-block:: python bar = Bar('Loading', fill='@', suffix='%(percent)d%%') This will produce a bar like the following: :: Loading |@@@@@@@@@@@@@ | 42% You can use a number of template arguments in ``message`` and ``suffix``: ========== ================================ Name Value ========== ================================ index current value max maximum value remaining max - index progress index / max percent progress * 100 avg simple moving average time per item (in seconds) elapsed elapsed time in seconds elapsed_td elapsed as a timedelta (useful for printing as a string) eta avg * remaining eta_td eta as a timedelta (useful for printing as a string) ========== ================================ Instead of passing all configuration options on instatiation, you can create your custom subclass: .. code-block:: python class FancyBar(Bar): message = 'Loading' fill = '*' suffix = '%(percent).1f%% - %(eta)ds' You can also override any of the arguments or create your own: .. code-block:: python class SlowBar(Bar): suffix = '%(remaining_hours)d hours remaining' @property def remaining_hours(self): return self.eta // 3600 Spinners ======== For actions with an unknown number of steps you can use a spinner: .. code-block:: python from progress.spinner import Spinner spinner = Spinner('Loading ') while state != 'FINISHED': # Do some work spinner.next() There are 5 predefined spinners: - ``Spinner`` - ``PieSpinner`` - ``MoonSpinner`` - ``LineSpinner`` - ``PixelSpinner`` Other ===== There are a number of other classes available too, please check the source or subclass one of them to create your own. License ======= progress is licensed under ISC ================================================ FILE: utils/progress/progress/__init__.py ================================================ # Copyright (c) 2012 Giorgos Verigakis # # Permission to use, copy, modify, and distribute this software for any # purpose with or without fee is hereby granted, provided that the above # copyright notice and this permission notice appear in all copies. # # THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES # WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF # MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR # ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES # WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN # ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF # OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. from __future__ import division from collections import deque from datetime import timedelta from math import ceil from sys import stderr from time import time __version__ = '1.3' class Infinite(object): file = stderr sma_window = 10 # Simple Moving Average window def __init__(self, *args, **kwargs): self.index = 0 self.start_ts = time() self.avg = 0 self._ts = self.start_ts self._xput = deque(maxlen=self.sma_window) for key, val in kwargs.items(): setattr(self, key, val) def __getitem__(self, key): if key.startswith('_'): return None return getattr(self, key, None) @property def elapsed(self): return int(time() - self.start_ts) @property def elapsed_td(self): return timedelta(seconds=self.elapsed) def update_avg(self, n, dt): if n > 0: self._xput.append(dt / n) self.avg = sum(self._xput) / len(self._xput) def update(self): pass def start(self): pass def finish(self): pass def next(self, n=1): now = time() dt = now - self._ts self.update_avg(n, dt) self._ts = now self.index = self.index + n self.update() def iter(self, it): try: for x in it: yield x self.next() finally: self.finish() class Progress(Infinite): def __init__(self, *args, **kwargs): super(Progress, self).__init__(*args, **kwargs) self.max = kwargs.get('max', 100) @property def eta(self): return int(ceil(self.avg * self.remaining)) @property def eta_td(self): return timedelta(seconds=self.eta) @property def percent(self): return self.progress * 100 @property def progress(self): return min(1, self.index / self.max) @property def remaining(self): return max(self.max - self.index, 0) def start(self): self.update() def goto(self, index): incr = index - self.index self.next(incr) def iter(self, it): try: self.max = len(it) except TypeError: pass try: for x in it: yield x self.next() finally: self.finish() ================================================ FILE: utils/progress/progress/bar.py ================================================ # -*- coding: utf-8 -*- # Copyright (c) 2012 Giorgos Verigakis # # Permission to use, copy, modify, and distribute this software for any # purpose with or without fee is hereby granted, provided that the above # copyright notice and this permission notice appear in all copies. # # THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES # WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF # MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR # ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES # WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN # ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF # OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. from __future__ import unicode_literals from . import Progress from .helpers import WritelnMixin class Bar(WritelnMixin, Progress): width = 32 message = '' suffix = '%(index)d/%(max)d' bar_prefix = ' |' bar_suffix = '| ' empty_fill = ' ' fill = '#' hide_cursor = True def update(self): filled_length = int(self.width * self.progress) empty_length = self.width - filled_length message = self.message % self bar = self.fill * filled_length empty = self.empty_fill * empty_length suffix = self.suffix % self line = ''.join([message, self.bar_prefix, bar, empty, self.bar_suffix, suffix]) self.writeln(line) class ChargingBar(Bar): suffix = '%(percent)d%%' bar_prefix = ' ' bar_suffix = ' ' empty_fill = '∙' fill = '█' class FillingSquaresBar(ChargingBar): empty_fill = '▢' fill = '▣' class FillingCirclesBar(ChargingBar): empty_fill = '◯' fill = '◉' class IncrementalBar(Bar): phases = (' ', '▏', '▎', '▍', '▌', '▋', '▊', '▉', '█') def update(self): nphases = len(self.phases) filled_len = self.width * self.progress nfull = int(filled_len) # Number of full chars phase = int((filled_len - nfull) * nphases) # Phase of last char nempty = self.width - nfull # Number of empty chars message = self.message % self bar = self.phases[-1] * nfull current = self.phases[phase] if phase > 0 else '' empty = self.empty_fill * max(0, nempty - len(current)) suffix = self.suffix % self line = ''.join([message, self.bar_prefix, bar, current, empty, self.bar_suffix, suffix]) self.writeln(line) class PixelBar(IncrementalBar): phases = ('⡀', '⡄', '⡆', '⡇', '⣇', '⣧', '⣷', '⣿') class ShadyBar(IncrementalBar): phases = (' ', '░', '▒', '▓', '█') ================================================ FILE: utils/progress/progress/counter.py ================================================ # -*- coding: utf-8 -*- # Copyright (c) 2012 Giorgos Verigakis # # Permission to use, copy, modify, and distribute this software for any # purpose with or without fee is hereby granted, provided that the above # copyright notice and this permission notice appear in all copies. # # THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES # WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF # MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR # ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES # WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN # ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF # OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. from __future__ import unicode_literals from . import Infinite, Progress from .helpers import WriteMixin class Counter(WriteMixin, Infinite): message = '' hide_cursor = True def update(self): self.write(str(self.index)) class Countdown(WriteMixin, Progress): hide_cursor = True def update(self): self.write(str(self.remaining)) class Stack(WriteMixin, Progress): phases = (' ', '▁', '▂', '▃', '▄', '▅', '▆', '▇', '█') hide_cursor = True def update(self): nphases = len(self.phases) i = min(nphases - 1, int(self.progress * nphases)) self.write(self.phases[i]) class Pie(Stack): phases = ('○', '◔', '◑', '◕', '●') ================================================ FILE: utils/progress/progress/helpers.py ================================================ # Copyright (c) 2012 Giorgos Verigakis # # Permission to use, copy, modify, and distribute this software for any # purpose with or without fee is hereby granted, provided that the above # copyright notice and this permission notice appear in all copies. # # THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES # WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF # MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR # ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES # WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN # ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF # OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. from __future__ import print_function HIDE_CURSOR = '\x1b[?25l' SHOW_CURSOR = '\x1b[?25h' class WriteMixin(object): hide_cursor = False def __init__(self, message=None, **kwargs): super(WriteMixin, self).__init__(**kwargs) self._width = 0 if message: self.message = message if self.file.isatty(): if self.hide_cursor: print(HIDE_CURSOR, end='', file=self.file) print(self.message, end='', file=self.file) self.file.flush() def write(self, s): if self.file.isatty(): b = '\b' * self._width c = s.ljust(self._width) print(b + c, end='', file=self.file) self._width = max(self._width, len(s)) self.file.flush() def finish(self): if self.file.isatty() and self.hide_cursor: print(SHOW_CURSOR, end='', file=self.file) class WritelnMixin(object): hide_cursor = False def __init__(self, message=None, **kwargs): super(WritelnMixin, self).__init__(**kwargs) if message: self.message = message if self.file.isatty() and self.hide_cursor: print(HIDE_CURSOR, end='', file=self.file) def clearln(self): if self.file.isatty(): print('\r\x1b[K', end='', file=self.file) def writeln(self, line): if self.file.isatty(): self.clearln() print(line, end='', file=self.file) self.file.flush() def finish(self): if self.file.isatty(): print(file=self.file) if self.hide_cursor: print(SHOW_CURSOR, end='', file=self.file) from signal import signal, SIGINT from sys import exit class SigIntMixin(object): """Registers a signal handler that calls finish on SIGINT""" def __init__(self, *args, **kwargs): super(SigIntMixin, self).__init__(*args, **kwargs) signal(SIGINT, self._sigint_handler) def _sigint_handler(self, signum, frame): self.finish() exit(0) ================================================ FILE: utils/progress/progress/spinner.py ================================================ # -*- coding: utf-8 -*- # Copyright (c) 2012 Giorgos Verigakis # # Permission to use, copy, modify, and distribute this software for any # purpose with or without fee is hereby granted, provided that the above # copyright notice and this permission notice appear in all copies. # # THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES # WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF # MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR # ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES # WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN # ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF # OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. from __future__ import unicode_literals from . import Infinite from .helpers import WriteMixin class Spinner(WriteMixin, Infinite): message = '' phases = ('-', '\\', '|', '/') hide_cursor = True def update(self): i = self.index % len(self.phases) self.write(self.phases[i]) class PieSpinner(Spinner): phases = ['◷', '◶', '◵', '◴'] class MoonSpinner(Spinner): phases = ['◑', '◒', '◐', '◓'] class LineSpinner(Spinner): phases = ['⎺', '⎻', '⎼', '⎽', '⎼', '⎻'] class PixelSpinner(Spinner): phases = ['⣾','⣷', '⣯', '⣟', '⡿', '⢿', '⣻', '⣽'] ================================================ FILE: utils/progress/setup.py ================================================ #!/usr/bin/env python from setuptools import setup import progress setup( name='progress', version=progress.__version__, description='Easy to use progress bars', long_description=open('README.rst').read(), author='Giorgos Verigakis', author_email='verigak@gmail.com', url='http://github.com/verigak/progress/', license='ISC', packages=['progress'], classifiers=[ 'Environment :: Console', 'Intended Audience :: Developers', 'License :: OSI Approved :: ISC License (ISCL)', 'Programming Language :: Python :: 2.6', 'Programming Language :: Python :: 2.7', 'Programming Language :: Python :: 3.3', 'Programming Language :: Python :: 3.4', 'Programming Language :: Python :: 3.5', 'Programming Language :: Python :: 3.6', ] ) ================================================ FILE: utils/progress/test_progress.py ================================================ #!/usr/bin/env python from __future__ import print_function import random import time from progress.bar import (Bar, ChargingBar, FillingSquaresBar, FillingCirclesBar, IncrementalBar, PixelBar, ShadyBar) from progress.spinner import (Spinner, PieSpinner, MoonSpinner, LineSpinner, PixelSpinner) from progress.counter import Counter, Countdown, Stack, Pie def sleep(): t = 0.01 t += t * random.uniform(-0.1, 0.1) # Add some variance time.sleep(t) for bar_cls in (Bar, ChargingBar, FillingSquaresBar, FillingCirclesBar): suffix = '%(index)d/%(max)d [%(elapsed)d / %(eta)d / %(eta_td)s]' bar = bar_cls(bar_cls.__name__, suffix=suffix) for i in bar.iter(range(200)): sleep() for bar_cls in (IncrementalBar, PixelBar, ShadyBar): suffix = '%(percent)d%% [%(elapsed_td)s / %(eta)d / %(eta_td)s]' bar = bar_cls(bar_cls.__name__, suffix=suffix) for i in bar.iter(range(200)): sleep() for spin in (Spinner, PieSpinner, MoonSpinner, LineSpinner, PixelSpinner): for i in spin(spin.__name__ + ' ').iter(range(100)): sleep() print() for singleton in (Counter, Countdown, Stack, Pie): for i in singleton(singleton.__name__ + ' ').iter(range(100)): sleep() print() bar = IncrementalBar('Random', suffix='%(index)d') for i in range(100): bar.goto(random.randint(0, 100)) sleep() bar.finish() ================================================ FILE: utils/radam.py ================================================ #!usr/bin/env python #-*- coding:utf-8 _*- """ @version: python3.6 @author: QLMX @contact: wenruichn@gmail.com @time: 2019-09-13 10:49 """ import math import torch from torch.optim.optimizer import Optimizer, required class RAdam(Optimizer): def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0): defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay) self.buffer = [[None, None, None] for ind in range(10)] super(RAdam, self).__init__(params, defaults) def __setstate__(self, state): super(RAdam, self).__setstate__(state) def step(self, closure=None): loss = None if closure is not None: loss = closure() for group in self.param_groups: for p in group['params']: if p.grad is None: continue grad = p.grad.data.float() if grad.is_sparse: raise RuntimeError('RAdam does not support sparse gradients') p_data_fp32 = p.data.float() state = self.state[p] if len(state) == 0: state['step'] = 0 state['exp_avg'] = torch.zeros_like(p_data_fp32) state['exp_avg_sq'] = torch.zeros_like(p_data_fp32) else: state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32) state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32) exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] beta1, beta2 = group['betas'] exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) exp_avg.mul_(beta1).add_(1 - beta1, grad) state['step'] += 1 buffered = self.buffer[int(state['step'] % 10)] if state['step'] == buffered[0]: N_sma, step_size = buffered[1], buffered[2] else: buffered[0] = state['step'] beta2_t = beta2 ** state['step'] N_sma_max = 2 / (1 - beta2) - 1 N_sma = N_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t) buffered[1] = N_sma # more conservative since it's an approximated value if N_sma >= 5: step_size = math.sqrt( (1 - beta2_t) * (N_sma - 4) / (N_sma_max - 4) * (N_sma - 2) / N_sma * N_sma_max / ( N_sma_max - 2)) / (1 - beta1 ** state['step']) else: step_size = 1.0 / (1 - beta1 ** state['step']) buffered[2] = step_size if group['weight_decay'] != 0: p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32) # more conservative since it's an approximated value if N_sma >= 5: denom = exp_avg_sq.sqrt().add_(group['eps']) p_data_fp32.addcdiv_(-step_size * group['lr'], exp_avg, denom) else: p_data_fp32.add_(-step_size * group['lr'], exp_avg) p.data.copy_(p_data_fp32) return loss class PlainRAdam(Optimizer): def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0): defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay) super(PlainRAdam, self).__init__(params, defaults) def __setstate__(self, state): super(PlainRAdam, self).__setstate__(state) def step(self, closure=None): loss = None if closure is not None: loss = closure() for group in self.param_groups: for p in group['params']: if p.grad is None: continue grad = p.grad.data.float() if grad.is_sparse: raise RuntimeError('RAdam does not support sparse gradients') p_data_fp32 = p.data.float() state = self.state[p] if len(state) == 0: state['step'] = 0 state['exp_avg'] = torch.zeros_like(p_data_fp32) state['exp_avg_sq'] = torch.zeros_like(p_data_fp32) else: state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32) state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32) exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] beta1, beta2 = group['betas'] exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) exp_avg.mul_(beta1).add_(1 - beta1, grad) state['step'] += 1 beta2_t = beta2 ** state['step'] N_sma_max = 2 / (1 - beta2) - 1 N_sma = N_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t) if group['weight_decay'] != 0: p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32) # more conservative since it's an approximated value if N_sma >= 5: step_size = group['lr'] * math.sqrt( (1 - beta2_t) * (N_sma - 4) / (N_sma_max - 4) * (N_sma - 2) / N_sma * N_sma_max / ( N_sma_max - 2)) / (1 - beta1 ** state['step']) denom = exp_avg_sq.sqrt().add_(group['eps']) p_data_fp32.addcdiv_(-step_size, exp_avg, denom) else: step_size = group['lr'] / (1 - beta1 ** state['step']) p_data_fp32.add_(-step_size, exp_avg) p.data.copy_(p_data_fp32) return loss class AdamW(Optimizer): def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, warmup=0): defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, warmup=warmup) super(AdamW, self).__init__(params, defaults) def __setstate__(self, state): super(AdamW, self).__setstate__(state) def step(self, closure=None): loss = None if closure is not None: loss = closure() for group in self.param_groups: for p in group['params']: if p.grad is None: continue grad = p.grad.data.float() if grad.is_sparse: raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead') p_data_fp32 = p.data.float() state = self.state[p] if len(state) == 0: state['step'] = 0 state['exp_avg'] = torch.zeros_like(p_data_fp32) state['exp_avg_sq'] = torch.zeros_like(p_data_fp32) else: state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32) state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32) exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] beta1, beta2 = group['betas'] state['step'] += 1 exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) exp_avg.mul_(beta1).add_(1 - beta1, grad) denom = exp_avg_sq.sqrt().add_(group['eps']) bias_correction1 = 1 - beta1 ** state['step'] bias_correction2 = 1 - beta2 ** state['step'] if group['warmup'] > state['step']: scheduled_lr = 1e-8 + state['step'] * group['lr'] / group['warmup'] else: scheduled_lr = group['lr'] step_size = scheduled_lr * math.sqrt(bias_correction2) / bias_correction1 if group['weight_decay'] != 0: p_data_fp32.add_(-group['weight_decay'] * scheduled_lr, p_data_fp32) p_data_fp32.addcdiv_(-step_size, exp_avg, denom) p.data.copy_(p_data_fp32) return loss ================================================ FILE: utils/utils.py ================================================ #!usr/bin/env python #-*- coding:utf-8 _*- """ @version: python3.6 @author: QLMX @contact: wenruichn@gmail.com @time: 2019-08-20 00:55 """ import torch.nn as nn import cv2 import torch device=torch.device("cuda") import os import urllib.parse as urlparse import requests import torch __all__ = ['GetEncoder', 'GetPreTrainedModel', 'load_pretrained', 'l2_norm'] #修改模型以进行特征提取 def GetEncoder(model): layerName,layer=list(model.named_children())[-1] exec("model."+layerName+"=nn.Linear(layer.in_features,layer.in_features)") exec("torch.nn.init.eye_(model."+layerName+".weight)") for param in model.parameters(): param.requires_grad=False return model,layer.in_features #修改模型以进行微调，n_ZeroChild和n_ZeroLayer用来设置参数固定层，当children为Sequential时使用n_ZeroLayer，可对其内部进行设置 def GetPreTrainedModel(model,n_Output,n_ZeroChild,n_ZeroLayer=None): for i,children in enumerate(model.children()): if i==n_ZeroChild: if n_ZeroLayer is not None: for j,layer in enumerate(children): if j==n_ZeroLayer: break for param in layer.parameters(): param.requires_grad=False break for param in children.parameters(): param.requires_grad=False layerName,layer=list(model.named_children())[-1] exec("model."+layerName+"=nn.Linear(layer.in_features,"+str(n_Output)+")") return model class StackNet2(nn.Module): def __init__(self,models,n_Target): super(StackNet,self).__init__() self.models=models n_Out=0 for i,(model,scale_In,n_ZeroChild,n_ZeroLayer) in enumerate(self.models): for j,children in enumerate(model.children()): if j==n_ZeroChild: if n_ZeroLayer is not None: for k,layer in enumerate(children): if k==n_ZeroLayer: break for param in layer.parameters(): param.requires_grad=False break layerName,layer=list(model.named_children())[-1] n_Out+=layer.in_features exec("model."+layerName+"=nn.Linear(layer.in_features,layer.in_features)") exec("torch.nn.init.eye_(model."+layerName+".weight)") exec("layer=model."+layerName) for param in layer.parameters(): param.requires_grad=False exec("self.model"+str(i)+"=model") self.fc=nn.Linear(n_Out,n_Target) def forward(self,x): feature=[] for model,scale_In,_,_ in self.models: feature.append(model(x)) feature=torch.cat(feature,dim=1) return self.fc(feature) def _download_file_from_google_drive(fid, dest): def _get_confirm_token(res): for k, v in res.cookies.items(): if k.startswith('download_warning'): return v return None def _save_response_content(res, dest): CHUNK_SIZE = 32768 with open(dest, "wb") as f: for chunk in res.iter_content(CHUNK_SIZE): if chunk: f.write(chunk) URL = "https://docs.google.com/uc?export=download" sess = requests.Session() res = sess.get(URL, params={'id': fid}, stream=True) token = _get_confirm_token(res) if token: params = {'id': fid, 'confirm': token} res = sess.get(URL, params=params, stream=True) _save_response_content(res, dest) def _load_url(url, dest): if os.path.isfile(dest) and os.path.exists(dest): return dest print('[INFO]: Downloading weights...') fid = urlparse.parse_qs(urlparse.urlparse(url).query)['id'][0] _download_file_from_google_drive(fid, dest) return dest def load_pretrained(m, meta, dest, pretrained=False): if pretrained: if len(meta) == 0: print('[INFO]: Pretrained model not available') return m if dest is None: dest = meta[0] else: dest = dest + '/' + meta[0] print(dest) m.load_state_dict(torch.load(_load_url(meta[1], dest))) return m def l2_norm(input,axis=1): norm = torch.norm(input,2,axis,True) output = torch.div(input, norm) return output ================================================ FILE: utils/visualize.py ================================================ import matplotlib.pyplot as plt import torch import torch.nn as nn import torchvision import torchvision.transforms as transforms import numpy as np from .misc import * __all__ = ['make_image', 'show_batch', 'show_mask', 'show_mask_single'] # functions to show an image def make_image(img, mean=(0,0,0), std=(1,1,1)): for i in range(0, 3): img[i] = img[i] * std[i] + mean[i] # unnormalize npimg = img.numpy() return np.transpose(npimg, (1, 2, 0)) def gauss(x,a,b,c): return torch.exp(-torch.pow(torch.add(x,-b),2).div(2*c*c)).mul(a) def colorize(x): ''' Converts a one-channel grayscale image to a color heatmap image ''' if x.dim() == 2: torch.unsqueeze(x, 0, out=x) if x.dim() == 3: cl = torch.zeros([3, x.size(1), x.size(2)]) cl[0] = gauss(x,.5,.6,.2) + gauss(x,1,.8,.3) cl[1] = gauss(x,1,.5,.3) cl[2] = gauss(x,1,.2,.3) cl[cl.gt(1)] = 1 elif x.dim() == 4: cl = torch.zeros([x.size(0), 3, x.size(2), x.size(3)]) cl[:,0,:,:] = gauss(x,.5,.6,.2) + gauss(x,1,.8,.3) cl[:,1,:,:] = gauss(x,1,.5,.3) cl[:,2,:,:] = gauss(x,1,.2,.3) return cl def show_batch(images, Mean=(2, 2, 2), Std=(0.5,0.5,0.5)): images = make_image(torchvision.utils.make_grid(images), Mean, Std) plt.imshow(images) plt.show() def show_mask_single(images, mask, Mean=(2, 2, 2), Std=(0.5,0.5,0.5)): im_size = images.size(2) # save for adding mask im_data = images.clone() for i in range(0, 3): im_data[:,i,:,:] = im_data[:,i,:,:] * Std[i] + Mean[i] # unnormalize images = make_image(torchvision.utils.make_grid(images), Mean, Std) plt.subplot(2, 1, 1) plt.imshow(images) plt.axis('off') # for b in range(mask.size(0)): # mask[b] = (mask[b] - mask[b].min())/(mask[b].max() - mask[b].min()) mask_size = mask.size(2) # print('Max %f Min %f' % (mask.max(), mask.min())) mask = (upsampling(mask, scale_factor=im_size/mask_size)) # mask = colorize(upsampling(mask, scale_factor=im_size/mask_size)) # for c in range(3): # mask[:,c,:,:] = (mask[:,c,:,:] - Mean[c])/Std[c] # print(mask.size()) mask = make_image(torchvision.utils.make_grid(0.3*im_data+0.7*mask.expand_as(im_data))) # mask = make_image(torchvision.utils.make_grid(0.3*im_data+0.7*mask), Mean, Std) plt.subplot(2, 1, 2) plt.imshow(mask) plt.axis('off') def show_mask(images, masklist, Mean=(2, 2, 2), Std=(0.5,0.5,0.5)): im_size = images.size(2) # save for adding mask im_data = images.clone() for i in range(0, 3): im_data[:,i,:,:] = im_data[:,i,:,:] * Std[i] + Mean[i] # unnormalize images = make_image(torchvision.utils.make_grid(images), Mean, Std) plt.subplot(1+len(masklist), 1, 1) plt.imshow(images) plt.axis('off') for i in range(len(masklist)): mask = masklist[i].data.cpu() # for b in range(mask.size(0)): # mask[b] = (mask[b] - mask[b].min())/(mask[b].max() - mask[b].min()) mask_size = mask.size(2) # print('Max %f Min %f' % (mask.max(), mask.min())) mask = (upsampling(mask, scale_factor=im_size/mask_size)) # mask = colorize(upsampling(mask, scale_factor=im_size/mask_size)) # for c in range(3): # mask[:,c,:,:] = (mask[:,c,:,:] - Mean[c])/Std[c] # print(mask.size()) mask = make_image(torchvision.utils.make_grid(0.3*im_data+0.7*mask.expand_as(im_data))) # mask = make_image(torchvision.utils.make_grid(0.3*im_data+0.7*mask), Mean, Std) plt.subplot(1+len(masklist), 1, i+2) plt.imshow(mask) plt.axis('off') # x = torch.zeros(1, 3, 3) # out = colorize(x) # out_im = make_image(out) # plt.imshow(out_im) # plt.show()