├── .gitignore ├── .idea ├── .gitignore ├── Pooling.iml ├── inspectionProfiles │ └── profiles_settings.xml ├── misc.xml └── modules.xml ├── Pooling ├── __pycache__ │ └── get_pooling.cpython-37.pyc ├── get_pooling.py └── pooling_method │ ├── AvgPool.py │ ├── LpPool.py │ ├── MaxPool.py │ ├── MixedPool.py │ ├── StochasticPool.py │ ├── __init__.py │ ├── __pycache__ │ ├── AvgPool.cpython-37.pyc │ ├── LpPool.cpython-37.pyc │ ├── MaxPool.cpython-37.pyc │ ├── MixedPool.cpython-37.pyc │ ├── StochasticPool.cpython-37.pyc │ ├── __init__.cpython-37.pyc │ └── lip_pooling.cpython-37.pyc │ ├── lip_pooling.py │ ├── median_pool.py │ └── softpool.py ├── README.md ├── __pycache__ ├── configs.cpython-37.pyc └── dataset_loader.cpython-37.pyc ├── configs.py ├── counter.py ├── dataset_loader.py ├── models ├── get_network.py └── net │ ├── __init__.py │ ├── __pycache__ │ ├── __init__.cpython-37.pyc │ ├── resnet.cpython-37.pyc │ └── vgg.cpython-37.pyc │ ├── densenet.py │ ├── resnet.py │ └── vgg.py ├── train.py └── utils ├── __pycache__ ├── train_engine.cpython-37.pyc └── util.cpython-37.pyc ├── train_engine.py └── util.py /.gitignore: -------------------------------------------------------------------------------- 1 | /.idea 2 | /log 3 | /ckpts 4 | /runs 5 | /__pycache__ 6 | -------------------------------------------------------------------------------- /.idea/.gitignore: -------------------------------------------------------------------------------- 1 | # Default ignored files 2 | /shelf/ 3 | /workspace.xml 4 | # Datasource local storage ignored files 5 | /../../../../:\Github\Pooling\.idea/dataSources/ 6 | /dataSources.local.xml 7 | # Editor-based HTTP Client requests 8 | /httpRequests/ 9 | -------------------------------------------------------------------------------- /.idea/Pooling.iml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | -------------------------------------------------------------------------------- /.idea/inspectionProfiles/profiles_settings.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 6 | -------------------------------------------------------------------------------- /.idea/misc.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /.idea/modules.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | -------------------------------------------------------------------------------- /Pooling/__pycache__/get_pooling.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rentainhe/pytorch-pooling/f050e157adde95b2721136812158b014e993eee8/Pooling/__pycache__/get_pooling.cpython-37.pyc -------------------------------------------------------------------------------- /Pooling/get_pooling.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from importlib import import_module 3 | import torch 4 | 5 | def get_pooling(__C): 6 | try: 7 | poolmethod_path = 'Pooling.pooling_method' 8 | pool_method = getattr(import_module(poolmethod_path), __C.pooling) 9 | return pool_method() 10 | except ImportError: 11 | print('the pool method name you entered is not supported yet') 12 | sys.exit() 13 | 14 | # class config: 15 | # def __init__(self): 16 | # self.pooling = 'stochastic' 17 | # 18 | # c = config() 19 | # pool = get_pooling(c) 20 | # p = pool(2,2) 21 | # x = torch.randn(1,128,4,4) 22 | # print(p(x).size()) -------------------------------------------------------------------------------- /Pooling/pooling_method/AvgPool.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | class AvgPool(nn.Module): 6 | def __init__(self, kernel_size, stride, padding=0): 7 | super(AvgPool,self).__init__() 8 | self.kernel_size = kernel_size 9 | self.stride = stride 10 | self.padding = padding 11 | 12 | def forward(self, x): 13 | x = F.avg_pool2d(x, kernel_size=self.kernel_size, stride=self.stride, padding=self.padding) 14 | return x 15 | 16 | def avg(): 17 | print("You are using Avg Pooling Method") 18 | return AvgPool -------------------------------------------------------------------------------- /Pooling/pooling_method/LpPool.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | class LpPool(nn.Module): 6 | def __init__(self, kernel_size, stride, padding=0, p=3): 7 | super(LpPool,self).__init__() 8 | self.kernel_size = kernel_size 9 | self.stride = stride 10 | assert self.kernel_size == self.stride, "we only support kernel_size=stride now" 11 | self.padding = padding 12 | self.p = p 13 | 14 | def lp_pool2d(self,input, p): 15 | return torch.mean(input ** p, (4, 5)) ** (1 / p) 16 | 17 | def forward(self, x): 18 | b,c,h,w = x.size() 19 | kh,kw = self.kernel_size,self.stride 20 | h = h//kh 21 | w = w//kw 22 | x = x.view(b,c,h,kh,w,kw).permute(0,1,2,4,3,5).contiguous() 23 | x = self.lp_pool2d(x,self.p) 24 | return x 25 | 26 | def Lp(): 27 | print("You are using Lp Pooling Method") 28 | return LpPool 29 | -------------------------------------------------------------------------------- /Pooling/pooling_method/MaxPool.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | class MaxPool(nn.Module): 6 | def __init__(self, kernel_size, stride, padding=0): 7 | super(MaxPool,self).__init__() 8 | self.kernel_size = kernel_size 9 | self.stride = stride 10 | self.padding = padding 11 | 12 | def forward(self, x): 13 | x = F.max_pool2d(x, kernel_size=self.kernel_size, stride=self.stride, padding=self.padding) 14 | return x 15 | 16 | def max(): 17 | print("You are using Max Pooling Method") 18 | return MaxPool -------------------------------------------------------------------------------- /Pooling/pooling_method/MixedPool.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | class mixedPool(nn.Module): 6 | def __init__(self,kernel_size, stride, padding=0, alpha=0.5): 7 | # nn.Module.__init__(self) 8 | super(mixedPool, self).__init__() 9 | alpha = torch.FloatTensor([alpha]) 10 | self.alpha = nn.Parameter(alpha) # nn.Parameter is special Variable 11 | self.kernel_size = kernel_size 12 | self.stride = stride 13 | self.padding = padding 14 | 15 | def forward(self, x): 16 | x = self.alpha * F.max_pool2d(x, self.kernel_size, self.stride, self.padding) + ( 17 | 1 - self.alpha) * F.avg_pool2d(x, self.kernel_size, self.stride, self.padding) 18 | return x 19 | 20 | def mixed(): 21 | print("You are using Mixed Pooling Method") 22 | return mixedPool -------------------------------------------------------------------------------- /Pooling/pooling_method/StochasticPool.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from torch.nn.modules.utils import _pair, _quadruple 5 | 6 | class StochasticPool2d(nn.Module): 7 | """ Stochastic 2D pooling, where prob(selecting index)~value of the activation 8 | IM_SIZE should be divisible by 2, not best implementation. 9 | based off: 10 | https://gist.github.com/rwightman/f2d3849281624be7c0f11c85c87c1598#file-median_pool-py-L5 11 | Args: 12 | kernel_size: size of pooling kernel, int or 2-tuple 13 | stride: pool stride, int or 2-tuple 14 | padding: pool padding, int or 4-tuple (l, r, t, b) as in pytorch F.pad 15 | same: override padding and enforce same padding, boolean 16 | """ 17 | 18 | def __init__(self, kernel_size=2, stride=2, padding=0, same=False): 19 | super(StochasticPool2d, self).__init__() 20 | self.kernel_size = _pair(kernel_size) # I don't know what this is but it works 21 | self.stride = _pair(stride) 22 | self.padding = _quadruple(padding) # convert to l, r, t, b 23 | self.same = same 24 | 25 | def _padding(self, x): 26 | if self.same: 27 | ih, iw = x.size()[2:] 28 | if ih % self.stride[0] == 0: 29 | ph = max(self.k[0] - self.stride[0], 0) 30 | else: 31 | ph = max(self.k[0] - (ih % self.stride[0]), 0) 32 | if iw % self.stride[1] == 0: 33 | pw = max(self.k[1] - self.stride[1], 0) 34 | else: 35 | pw = max(self.k[1] - (iw % self.stride[1]), 0) 36 | pl = pw // 2 37 | pr = pw - pl 38 | pt = ph // 2 39 | pb = ph - pt 40 | padding = (pl, pr, pt, pb) 41 | else: 42 | padding = self.padding 43 | return padding 44 | 45 | def forward(self, x): 46 | # because multinomial likes to fail on GPU when all values are equal 47 | # Try randomly sampling without calling the get_random function a million times 48 | init_size = x.shape 49 | 50 | # x = F.pad(x, self._padding(x), mode='reflect') 51 | x = x.unfold(2, self.kernel_size[0], self.stride[0]).unfold(3, self.kernel_size[1], self.stride[1]) 52 | x = x.contiguous().view(-1, 4) 53 | idx = torch.randint(0, x.shape[1], size=(x.shape[0],)).type(torch.cuda.LongTensor) 54 | x = x.contiguous().view(-1) 55 | x = torch.take(x, idx) 56 | x = x.contiguous().view(init_size[0], init_size[1], int(init_size[2] / 2), int(init_size[3] / 2)) 57 | return x 58 | 59 | def stochastic(): 60 | print("You are using Stochatic Pooling Method") 61 | return StochasticPool2d -------------------------------------------------------------------------------- /Pooling/pooling_method/__init__.py: -------------------------------------------------------------------------------- 1 | from Pooling.pooling_method.AvgPool import * 2 | from Pooling.pooling_method.MaxPool import * 3 | from Pooling.pooling_method.MixedPool import * 4 | from Pooling.pooling_method.LpPool import * 5 | from Pooling.pooling_method.lip_pooling import * 6 | from Pooling.pooling_method.StochasticPool import * 7 | from Pooling.pooling_method.softpool import soft -------------------------------------------------------------------------------- /Pooling/pooling_method/__pycache__/AvgPool.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rentainhe/pytorch-pooling/f050e157adde95b2721136812158b014e993eee8/Pooling/pooling_method/__pycache__/AvgPool.cpython-37.pyc -------------------------------------------------------------------------------- /Pooling/pooling_method/__pycache__/LpPool.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rentainhe/pytorch-pooling/f050e157adde95b2721136812158b014e993eee8/Pooling/pooling_method/__pycache__/LpPool.cpython-37.pyc -------------------------------------------------------------------------------- /Pooling/pooling_method/__pycache__/MaxPool.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rentainhe/pytorch-pooling/f050e157adde95b2721136812158b014e993eee8/Pooling/pooling_method/__pycache__/MaxPool.cpython-37.pyc -------------------------------------------------------------------------------- /Pooling/pooling_method/__pycache__/MixedPool.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rentainhe/pytorch-pooling/f050e157adde95b2721136812158b014e993eee8/Pooling/pooling_method/__pycache__/MixedPool.cpython-37.pyc -------------------------------------------------------------------------------- /Pooling/pooling_method/__pycache__/StochasticPool.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rentainhe/pytorch-pooling/f050e157adde95b2721136812158b014e993eee8/Pooling/pooling_method/__pycache__/StochasticPool.cpython-37.pyc -------------------------------------------------------------------------------- /Pooling/pooling_method/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rentainhe/pytorch-pooling/f050e157adde95b2721136812158b014e993eee8/Pooling/pooling_method/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /Pooling/pooling_method/__pycache__/lip_pooling.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rentainhe/pytorch-pooling/f050e157adde95b2721136812158b014e993eee8/Pooling/pooling_method/__pycache__/lip_pooling.cpython-37.pyc -------------------------------------------------------------------------------- /Pooling/pooling_method/lip_pooling.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from collections import OrderedDict 5 | 6 | def lip2d(x, logit, kernel=2, stride=2, padding=0): 7 | weight = logit.exp() 8 | return F.avg_pool2d(x*weight, kernel, stride, padding)/F.avg_pool2d(weight, kernel, stride, padding) 9 | 10 | class SoftGate(nn.Module): 11 | def __init__(self, coeff=12): 12 | super(SoftGate, self).__init__() 13 | self.coeff = coeff 14 | 15 | def forward(self, x): 16 | return torch.sigmoid(x).mul(self.coeff) 17 | 18 | class SimplifiedLIP(nn.Module): 19 | def __init__(self, channels): 20 | super(SimplifiedLIP, self).__init__() 21 | 22 | rp = channels 23 | # nn.Sequential + OrderedDict 可以为每一个层赋予一个名字，来替代掉本身的 index 命名，但是只能通过index访问 24 | self.logit = nn.Sequential( 25 | OrderedDict(( 26 | ('conv', nn.Conv2d(channels, channels, 3, padding=1, bias=False)), 27 | ('bn', nn.InstanceNorm2d(channels, affine=True)), 28 | ('gate', SoftGate()), 29 | )) 30 | ) 31 | 32 | def init_layer(self): 33 | self.logit[0].weight.data.fill_(0.0) 34 | 35 | def forward(self, x): 36 | frac = lip2d(x, self.logit(x)) 37 | return frac 38 | 39 | def lip(): 40 | print("You are using Lip Pooling Method") 41 | return SimplifiedLIP 42 | -------------------------------------------------------------------------------- /Pooling/pooling_method/median_pool.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | from torch.nn.modules.utils import _pair, _quadruple 6 | 7 | 8 | class MedianPool2d(nn.Module): 9 | """ Median pool (usable as median filter when stride=1) module. 10 | Args: 11 | kernel_size: size of pooling kernel, int or 2-tuple 12 | stride: pool stride, int or 2-tuple 13 | padding: pool padding, int or 4-tuple (l, r, t, b) as in pytorch F.pad 14 | same: override padding and enforce same padding, boolean 15 | """ 16 | def __init__(self, kernel_size=3, stride=1, padding=0, same=False): 17 | super(MedianPool2d, self).__init__() 18 | self.k = _pair(kernel_size) 19 | self.stride = _pair(stride) 20 | self.padding = _quadruple(padding) # convert to l, r, t, b 21 | self.same = same 22 | 23 | def _padding(self, x): 24 | if self.same: 25 | ih, iw = x.size()[2:] 26 | if ih % self.stride[0] == 0: 27 | ph = max(self.k[0] - self.stride[0], 0) 28 | else: 29 | ph = max(self.k[0] - (ih % self.stride[0]), 0) 30 | if iw % self.stride[1] == 0: 31 | pw = max(self.k[1] - self.stride[1], 0) 32 | else: 33 | pw = max(self.k[1] - (iw % self.stride[1]), 0) 34 | pl = pw // 2 35 | pr = pw - pl 36 | pt = ph // 2 37 | pb = ph - pt 38 | padding = (pl, pr, pt, pb) 39 | else: 40 | padding = self.padding 41 | return padding 42 | 43 | def forward(self, x): 44 | x = F.pad(x, self._padding(x), mode='reflect') 45 | x = x.unfold(2, self.k[0], self.stride[0]).unfold(3, self.k[1], self.stride[1]) 46 | x = x.contiguous().view(x.size()[:4] + (-1,)).median(dim=-1)[0] 47 | return x -------------------------------------------------------------------------------- /Pooling/pooling_method/softpool.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from torch.nn.modules.utils import _triple, _pair, _single 5 | 6 | class SoftPool2d(nn.Module): 7 | def __init__(self, kernel_size, stride, padding=0): 8 | super(SoftPool2d,self).__init__() 9 | self.kernel_size = kernel_size 10 | self.stride = stride 11 | self.padding = padding 12 | 13 | def forward(self, x): 14 | x = self.soft_pool2d(x, kernel_size=self.kernel_size, stride=self.stride) 15 | return x 16 | 17 | def soft_pool2d(self, x, kernel_size=2, stride=None, force_inplace=False): 18 | kernel_size = _pair(kernel_size) 19 | if stride is None: 20 | stride = kernel_size 21 | else: 22 | stride = _pair(stride) 23 | _, c, h, w = x.size() 24 | e_x = torch.sum(torch.exp(x),dim=1,keepdim=True) 25 | return F.avg_pool2d(x.mul(e_x), kernel_size, stride=stride).mul_(sum(kernel_size)).div_(F.avg_pool2d(e_x, kernel_size, stride=stride).mul_(sum(kernel_size))) 26 | 27 | def soft(): 28 | print("You are using Soft Pooling Method") 29 | return SoftPool2d -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Pooling 2 | This is a collection of different pooling methods used in `image classification` `segmentation` `detection` 3 | 4 | ## Features 5 | * Multi-GPU support 6 | * Easy and Useful Training log file 7 | * Easy to test different pooling method on classification task 8 | 9 | ## Requirements 10 | * python3.6 11 | * pytorch1.6.0 + cuda10.1 12 | * tensorboard 2.3.0 13 | 14 | ## Installation 15 | * clone 16 | ``` 17 | git clone https://github.com/rentainhe/pytorch-pooling.git 18 | ``` 19 | * make data directory for cifar100 20 | ```bash 21 | $ cd pytorch-pooling 22 | $ mkdir data 23 | ``` 24 | 25 | ## Usage 26 | 27 | ### 1. enter directory 28 | ```bash 29 | $ cd pytorch-pooling 30 | ``` 31 | 32 | ### 2. dataset 33 | * Only support cifar100 now (Will support Imagenet Later) 34 | * Using cifar100 dataset from torchvision since it's more convinient 35 | 36 | ### 3. run tensorboard 37 | Install tensorboard 38 | ```bash 39 | $ pip install tensorboard 40 | Run tensorboard 41 | $ tensorboard --logdir runs --port 6006 --host localhost 42 | ``` 43 | 44 | ### 4. training 45 | Our base backbone is `vgg16` with `batch_normalization` 46 | ```bash 47 | $ python3 train.py --run train --name test --pooling max 48 | ``` 49 | 50 | - ```--run={'train','test','visual'}``` to set the mode to be executed 51 | 52 | - ```--name=str``` to set the name of this training 53 | 54 | - ```--pooling=str```, e.g, `--pooling='max'` to set the __pooling method__ in `vgg16` to be `max_pool2d` 55 | 56 | - ```--gpu=str```, e.g, `--gpu='1'` to set the specified GPU for training 57 | 58 | The supported pooling args are 59 | ``` 60 | max pooling 61 | average pooling 62 | mixed pooling 63 | Lp pooling 64 | lip pooling 65 | soft pooling 66 | ``` 67 | 68 | ### 5. Add a new pooling method 69 | You should add a new pooling method `pool.py` in `"/Pooling/pooling_method"` and update the `__init__.py` file 70 | 71 | ### 6. Addition 72 | - lip pooling: the backbone in original paper is `ResNet`, But I use `VggNet` in this repo, so there might be something wrong with the accuracy 73 | 74 | ## Results 75 | The result I can get from this repo, I train every model with the same hyperparam and I don't use any tricks in this repo. 76 | 77 | |dataset|backbone|pooling|acc|epoch(lr = 0.1)|epoch(lr = 0.02)|epoch(lr = 0.004)|epoch(lr = 0.0008)|total epoch| 78 | |:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---: 79 | |cifar100|vgg16_bn|max|70.89%|60|60|40|40|200| 80 | |cifar100|vgg16_bn|avg|70.56%|60|60|40|40|200| 81 | |cifar100|vgg16_bn|mixed|71.19%|60|60|40|40|200| 82 | |cifar100|vgg16_bn|Lp(p=2)|70.65%|60|60|40|40|200| 83 | |cifar100|vgg16_bn|Lp(p=3)|70.67%|60|60|40|40|200| 84 | |cifar100|vgg16_bn|lip|71.23%|60|60|40|40|200| 85 | |cifar100|vgg16_bn|softpool|__71.39%__|60|60|40|40|200| 86 | 87 | ## Implementated Pooling 88 | 89 | - mixed pooling [Mixed pooling for convolutional neural networks](https://rd.springer.com/chapter/10.1007/978-3-319-11740-9_34) 90 | - Lp pooling [Convolutional neural networks applied to house numbers digit 91 | classification](https://arxiv.org/abs/1204.3968) 92 | - lip pooling [LIP: Local Importance-based Pooling](https://arxiv.org/abs/1908.04156) 93 | - soft pooling [Refining activation downsampling with SoftPool](https://arxiv.org/abs/2101.00440) -------------------------------------------------------------------------------- /__pycache__/configs.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rentainhe/pytorch-pooling/f050e157adde95b2721136812158b014e993eee8/__pycache__/configs.cpython-37.pyc -------------------------------------------------------------------------------- /__pycache__/dataset_loader.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rentainhe/pytorch-pooling/f050e157adde95b2721136812158b014e993eee8/__pycache__/dataset_loader.cpython-37.pyc -------------------------------------------------------------------------------- /configs.py: -------------------------------------------------------------------------------- 1 | from types import MethodType 2 | from datetime import datetime 3 | import os 4 | import torch 5 | import numpy as np 6 | import random 7 | 8 | class Configs(): 9 | def __init__(self): 10 | self.epoch = 200 11 | self.milestones = [60, 120, 160] 12 | self.save_epoch = 20 13 | self.gpu='0' 14 | self.batch_size=64 15 | self.mean = (0.5070751592371323, 0.48654887331495095, 0.4409178433670343) 16 | self.std = (0.2673342858792401, 0.2564384629170883, 0.27615047132568404) 17 | self.num_workers = 8 18 | self.pin_memory = True 19 | self.eval_every_epoch = True 20 | self.time = datetime.now().strftime('%A_%d_%B_%Y_%Hh_%Mm_%Ss') 21 | self.seed = random.randint(0, 9999999) 22 | self.version = str(self.seed) 23 | self.batch_size = 128 24 | self.gradient_accumulation_steps = 1 25 | self.ckpts_dir = 'ckpts' 26 | self.result_log_dir = 'log' 27 | self.tensorboard_log_dir = 'runs' 28 | 29 | def parse_to_dict(self, args): 30 | args_dict = {} 31 | for arg in dir(args): 32 | if not arg.startswith('_') and not isinstance(getattr(args, arg), MethodType): 33 | if getattr(args, arg) is not None: 34 | args_dict[arg] = getattr(args, arg) 35 | return args_dict 36 | 37 | def add_args(self, args_dict): 38 | for arg in args_dict: 39 | setattr(self, arg, args_dict[arg]) 40 | 41 | def training_init(self): 42 | os.environ['CUDA_VISIBLE_DEVICES'] = self.gpu 43 | self.n_gpu = len(self.gpu.split(',')) 44 | self.devices = [_ for _ in range(self.n_gpu)] 45 | torch.set_num_threads(2) 46 | 47 | # fix seed 48 | torch.manual_seed(self.seed) 49 | if self.n_gpu < 2: 50 | torch.cuda.manual_seed(self.seed) 51 | else: 52 | torch.cuda.manual_seed_all(self.seed) 53 | torch.backends.cudnn.deterministic = True 54 | np.random.seed(self.seed) 55 | random.seed(self.seed) 56 | 57 | # Gradient accumulate setup 58 | assert self.batch_size % self.gradient_accumulation_steps == 0 59 | self.sub_batch_size = int(self.batch_size / self.gradient_accumulation_steps) 60 | self.eval_batch_size = int(self.sub_batch_size / 2) 61 | 62 | def path_init(self): 63 | for attr in dir(self): 64 | if 'dir' in attr and not attr.startswith('__'): 65 | if getattr(self,attr) not in os.listdir('./'): 66 | os.makedirs(getattr(self, attr)) 67 | 68 | 69 | def __str__(self): 70 | # print Hyper Parameters 71 | settings_str = '' 72 | for attr in dir(self): 73 | if not 'np' in attr and not 'random' in attr and not attr.startswith('__') and not isinstance(getattr(self, attr), MethodType): 74 | settings_str += '{ %-17s }->' % attr + str(getattr(self, attr)) + '\n' 75 | return settings_str 76 | 77 | configs = Configs() 78 | 79 | 80 | 81 | 82 | -------------------------------------------------------------------------------- /counter.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from ptflops import get_model_complexity_info 4 | from models.get_network import get_network 5 | import argparse 6 | from models.net.vgg import vgg16_bn 7 | 8 | 9 | def count_parameters(net): 10 | params = sum([param.nelement() for param in net.parameters() if param.requires_grad]) 11 | print("Params: %f M" % (params / 1000000)) 12 | 13 | 14 | def parse_args(): 15 | parser = argparse.ArgumentParser(description='MAC_Classification Args') 16 | parser.add_argument('--pooling', type=str, 17 | choices=[ 18 | 'max', 19 | 'avg', 20 | 'mixed', 21 | 'Lp', 22 | 'lip', 23 | 'stochastic', 24 | 'soft' 25 | ], default='max', help='choose one pooling method to use', required=True) 26 | parser.add_argument('--gpu', type=int, default=0, help="choose a gpu for testing") 27 | args = parser.parse_args() 28 | return args 29 | 30 | 31 | if __name__ == '__main__': 32 | 33 | device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 34 | args = parse_args() 35 | 36 | if device == 'cpu': 37 | net = vgg16_bn(args) 38 | count_parameters(net) 39 | else: 40 | with torch.cuda.device(args.gpu): 41 | net = vgg16_bn(args) 42 | macs, params = get_model_complexity_info(net, (3, 224, 224), as_strings=True, 43 | print_per_layer_stat=True, verbose=True) 44 | print('{:<30} {:<8}'.format('Computational complexity: ', macs)) 45 | print('{:<30} {:<8}'.format('Number of parameters: ', params)) 46 | -------------------------------------------------------------------------------- /dataset_loader.py: -------------------------------------------------------------------------------- 1 | from torchvision import transforms, datasets 2 | from torch.utils.data import DataLoader, RandomSampler, SequentialSampler 3 | 4 | def get_train_loader(configs): 5 | # data transforms 6 | transform_train = transforms.Compose([ 7 | transforms.RandomCrop(configs.img_size, padding=4), 8 | transforms.ToTensor(), 9 | transforms.Normalize(configs.mean, configs.std) 10 | ]) 11 | trainset = datasets.CIFAR100(root='./data/cifar100', 12 | train=True, 13 | download=True, 14 | transform=transform_train) 15 | # get dataloader 16 | train_sampler = RandomSampler(trainset) 17 | train_loader = DataLoader(trainset, 18 | sampler=train_sampler, 19 | batch_size=configs.batch_size, 20 | num_workers=configs.num_workers, 21 | pin_memory=configs.pin_memory) 22 | return train_loader 23 | 24 | def get_test_loader(configs): 25 | # data transforms 26 | transform_test = transforms.Compose([ 27 | transforms.Resize((configs.img_size, configs.img_size)), 28 | transforms.ToTensor(), 29 | transforms.Normalize(configs.mean, configs.std) 30 | ]) 31 | testset = datasets.CIFAR100(root='./data/cifar100', 32 | train=False, 33 | download=True, 34 | transform=transform_test) 35 | 36 | # get dataloader 37 | test_sampler = SequentialSampler(testset) 38 | test_loader = DataLoader(testset, 39 | sampler=test_sampler, 40 | batch_size=configs.eval_batch_size, 41 | num_workers=configs.num_workers, 42 | pin_memory=configs.pin_memory) 43 | return test_loader 44 | -------------------------------------------------------------------------------- /models/get_network.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import importlib 4 | from importlib import import_module 5 | 6 | def get_network(__C): 7 | try: 8 | model_path = 'models.net' 9 | net = getattr(import_module(model_path),__C.model) 10 | return net() 11 | except ImportError: 12 | print('the network name you have entered is not supported yet') 13 | sys.exit() 14 | -------------------------------------------------------------------------------- /models/net/__init__.py: -------------------------------------------------------------------------------- 1 | from models.net.resnet import * 2 | from models.net.vgg import * -------------------------------------------------------------------------------- /models/net/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rentainhe/pytorch-pooling/f050e157adde95b2721136812158b014e993eee8/models/net/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /models/net/__pycache__/resnet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rentainhe/pytorch-pooling/f050e157adde95b2721136812158b014e993eee8/models/net/__pycache__/resnet.cpython-37.pyc -------------------------------------------------------------------------------- /models/net/__pycache__/vgg.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rentainhe/pytorch-pooling/f050e157adde95b2721136812158b014e993eee8/models/net/__pycache__/vgg.cpython-37.pyc -------------------------------------------------------------------------------- /models/net/densenet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | import math 5 | 6 | 7 | __all__ = ['densenet'] 8 | 9 | 10 | from torch.autograd import Variable 11 | 12 | class Bottleneck(nn.Module): 13 | def __init__(self, inplanes, expansion=4, growthRate=12, dropRate=0): 14 | super(Bottleneck, self).__init__() 15 | planes = expansion * growthRate 16 | self.bn1 = nn.BatchNorm2d(inplanes) 17 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) 18 | self.bn2 = nn.BatchNorm2d(planes) 19 | self.conv2 = nn.Conv2d(planes, growthRate, kernel_size=3, 20 | padding=1, bias=False) 21 | self.relu = nn.ReLU(inplace=True) 22 | self.dropRate = dropRate 23 | 24 | def forward(self, x): 25 | out = self.bn1(x) 26 | out = self.relu(out) 27 | out = self.conv1(out) 28 | out = self.bn2(out) 29 | out = self.relu(out) 30 | out = self.conv2(out) 31 | if self.dropRate > 0: 32 | out = F.dropout(out, p=self.dropRate, training=self.training) 33 | 34 | out = torch.cat((x, out), 1) 35 | 36 | return out 37 | 38 | 39 | class BasicBlock(nn.Module): 40 | def __init__(self, inplanes, expansion=1, growthRate=12, dropRate=0): 41 | super(BasicBlock, self).__init__() 42 | planes = expansion * growthRate 43 | self.bn1 = nn.BatchNorm2d(inplanes) 44 | self.conv1 = nn.Conv2d(inplanes, growthRate, kernel_size=3, 45 | padding=1, bias=False) 46 | self.relu = nn.ReLU(inplace=True) 47 | self.dropRate = dropRate 48 | 49 | def forward(self, x): 50 | out = self.bn1(x) 51 | out = self.relu(out) 52 | out = self.conv1(out) 53 | if self.dropRate > 0: 54 | out = F.dropout(out, p=self.dropRate, training=self.training) 55 | 56 | out = torch.cat((x, out), 1) 57 | 58 | return out 59 | 60 | 61 | class Transition(nn.Module): 62 | def __init__(self, inplanes, outplanes): 63 | super(Transition, self).__init__() 64 | self.bn1 = nn.BatchNorm2d(inplanes) 65 | self.conv1 = nn.Conv2d(inplanes, outplanes, kernel_size=1, 66 | bias=False) 67 | self.relu = nn.ReLU(inplace=True) 68 | 69 | def forward(self, x): 70 | out = self.bn1(x) 71 | out = self.relu(out) 72 | out = self.conv1(out) 73 | out = F.avg_pool2d(out, 2) 74 | return out 75 | 76 | 77 | class DenseNet(nn.Module): 78 | 79 | def __init__(self, depth=22, block=Bottleneck, 80 | dropRate=0, num_classes=10, growthRate=12, compressionRate=2): 81 | super(DenseNet, self).__init__() 82 | 83 | assert (depth - 4) % 3 == 0, 'depth should be 3n+4' 84 | n = (depth - 4) / 3 if block == BasicBlock else (depth - 4) // 6 85 | 86 | self.growthRate = growthRate 87 | self.dropRate = dropRate 88 | 89 | # self.inplanes is a global variable used across multiple 90 | # helper functions 91 | self.inplanes = growthRate * 2 92 | self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, padding=1, 93 | bias=False) 94 | self.dense1 = self._make_denseblock(block, n) 95 | self.trans1 = self._make_transition(compressionRate) 96 | self.dense2 = self._make_denseblock(block, n) 97 | self.trans2 = self._make_transition(compressionRate) 98 | self.dense3 = self._make_denseblock(block, n) 99 | self.bn = nn.BatchNorm2d(self.inplanes) 100 | self.relu = nn.ReLU(inplace=True) 101 | self.avgpool = nn.AvgPool2d(8) 102 | self.fc = nn.Linear(self.inplanes, num_classes) 103 | 104 | # Weight initialization 105 | for m in self.modules(): 106 | if isinstance(m, nn.Conv2d): 107 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 108 | m.weight.data.normal_(0, math.sqrt(2. / n)) 109 | elif isinstance(m, nn.BatchNorm2d): 110 | m.weight.data.fill_(1) 111 | m.bias.data.zero_() 112 | 113 | def _make_denseblock(self, block, blocks): 114 | layers = [] 115 | for i in range(blocks): 116 | # Currently we fix the expansion ratio as the default value 117 | layers.append(block(self.inplanes, growthRate=self.growthRate, dropRate=self.dropRate)) 118 | self.inplanes += self.growthRate 119 | 120 | return nn.Sequential(*layers) 121 | 122 | def _make_transition(self, compressionRate): 123 | inplanes = self.inplanes 124 | outplanes = int(math.floor(self.inplanes // compressionRate)) 125 | self.inplanes = outplanes 126 | return Transition(inplanes, outplanes) 127 | 128 | 129 | def forward(self, x): 130 | x = self.conv1(x) 131 | 132 | x = self.trans1(self.dense1(x)) 133 | x = self.trans2(self.dense2(x)) 134 | x = self.dense3(x) 135 | x = self.bn(x) 136 | x = self.relu(x) 137 | 138 | x = self.avgpool(x) 139 | x = x.view(x.size(0), -1) 140 | x = self.fc(x) 141 | 142 | return x 143 | 144 | 145 | def densenet(**kwargs): 146 | """ 147 | Constructs a ResNet model. 148 | """ 149 | return DenseNet(**kwargs) -------------------------------------------------------------------------------- /models/net/resnet.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | class BasicBlock(nn.Module): 4 | # Residual block for resnet18 and resnet 34 5 | expansion = 1 # 维度拓展 6 | def __init__(self, in_channels, out_channels, stride=1): 7 | super().__init__() 8 | # residual function 9 | self.residual_function = nn.Sequential( 10 | nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False), 11 | nn.BatchNorm2d(out_channels), 12 | nn.ReLU(inplace=True), 13 | nn.Conv2d(out_channels, out_channels * BasicBlock.expansion, kernel_size=3, padding=1, bias=False), 14 | nn.BatchNorm2d(out_channels * BasicBlock.expansion) # 注意 expansion 的用法 15 | ) 16 | self.shortcut = nn.Sequential() 17 | 18 | # the shortcut output dimension is not the same with residual function 19 | # use 1*1 convolution to match the dimension 20 | if stride!=1 or in_channels!=BasicBlock.expansion * out_channels: 21 | self.shortcut = nn.Sequential( 22 | nn.Conv2d(in_channels, out_channels * BasicBlock.expansion, kernel_size=1, stride=stride, bias=False), 23 | nn.BatchNorm2d(out_channels * BasicBlock.expansion) 24 | ) 25 | 26 | def forward(self, x): 27 | return nn.ReLU(inplace=True)(self.residual_function(x) + self.shortcut(x)) 28 | 29 | class BottleNeck(nn.Module): 30 | # Residual block for resnet over 50 layers 31 | # 瓶颈层，因为用了 1*1 的卷积，比较方便改变网络的大小，先降维后升维，将高频噪声消除，并且减少计算量，在深层网络上常用 32 | expansion=4 33 | def __init__(self, in_channels, out_channels, stride=1): 34 | super().__init__() 35 | self.residual_function = nn.Sequential( 36 | nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False), 37 | nn.BatchNorm2d(out_channels), 38 | nn.ReLU(inplace=True), 39 | nn.Conv2d(out_channels, out_channels, stride=stride, kernel_size=3, padding=1, bias=False), 40 | nn.BatchNorm2d(out_channels), 41 | nn.ReLU(inplace=True), 42 | nn.Conv2d(out_channels, out_channels * BottleNeck.expansion, kernel_size=1, bias=False), 43 | nn.BatchNorm2d(out_channels * BottleNeck.expansion) 44 | ) 45 | self.shortcut = nn.Sequential() 46 | 47 | # 如果 shortcut 的维度和 residual function 的维度不相等的话，需要用 1*1 的卷积改变一下维度大小 48 | if stride!=1 or in_channels!=out_channels*BottleNeck.expansion: 49 | self.shortcut = nn.Sequential( 50 | nn.Conv2d(in_channels, out_channels * BottleNeck.expansion, stride=stride, kernel_size=1, bias=False), 51 | nn.BatchNorm2d(out_channels * BottleNeck.expansion) 52 | ) 53 | 54 | def forward(self, x): 55 | return nn.ReLU(inplace=True)(self.residual_function(x) + self.shortcut(x)) 56 | 57 | class ResNet(nn.Module): 58 | def __init__(self, block, num_block, num_classes=100): 59 | super().__init__() 60 | self.in_channels = 64 61 | 62 | self.conv1 = nn.Sequential( 63 | nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=False), 64 | nn.BatchNorm2d(64), 65 | nn.ReLU(inplace=True) 66 | ) 67 | 68 | self.conv2_x = self._make_layer(block, 64, num_block[0], 1) 69 | self.conv3_x = self._make_layer(block, 128, num_block[1], 2) 70 | self.conv4_x = self._make_layer(block, 256, num_block[2], 2) 71 | self.conv5_x = self._make_layer(block, 512, num_block[3], 2) 72 | self.avg_pool = nn.AdaptiveAvgPool2d((1,1)) 73 | self.fc = nn.Linear(512*block.expansion, num_classes) 74 | 75 | def _make_layer(self, block, out_channels, num_blocks, stride): 76 | """make resnet layers(by layer i didn't mean this 'layer' was the 77 | same as a neuron network layer, ex. conv layer), one layer may 78 | contain more than one residual block 79 | 80 | Args: 81 | block: block type, basic block or bottle neck block 82 | out_channels: output depth channel number of this layer 83 | num_blocks: how many blocks per layer 84 | stride: the stride of the first block of this layer 85 | 86 | Return: 87 | return a resnet layer 88 | """ 89 | strides = [stride] + [1]*(num_blocks-1) 90 | layers = [] 91 | for stride in strides: 92 | layers.append(block(self.in_channels, out_channels, stride)) 93 | self.in_channels = out_channels * block.expansion 94 | 95 | return nn.Sequential(*layers) 96 | 97 | def forward(self, x): 98 | output = self.conv1(x) 99 | output = self.conv2_x(output) 100 | output = self.conv3_x(output) 101 | output = self.conv4_x(output) 102 | output = self.conv5_x(output) 103 | output = self.avg_pool(output) 104 | output = output.view(output.size(0),-1) 105 | output = self.fc(output) 106 | 107 | return output 108 | 109 | 110 | def resnet18(): 111 | # return a resnet18 object 112 | return ResNet(BasicBlock, [2,2,2,2]) 113 | 114 | 115 | def resnet34(): 116 | # return a resnet34 object 117 | return ResNet(BasicBlock, [3,4,6,3]) 118 | 119 | 120 | def resnet50(): 121 | # return a resnet50 object 122 | return ResNet(BottleNeck, [3,4,6,3]) -------------------------------------------------------------------------------- /models/net/vgg.py: -------------------------------------------------------------------------------- 1 | """vgg in pytorch 2 | [1] Karen Simonyan, Andrew Zisserman 3 | 4 | Very Deep Convolutional Networks for Large-Scale Image Recognition. 5 | https://arxiv.org/abs/1409.1556v6 6 | """ 7 | '''VGG11/13/16/19 in Pytorch.''' 8 | 9 | import torch 10 | import torch.nn as nn 11 | from Pooling.get_pooling import get_pooling 12 | 13 | 14 | cfg = { 15 | 'A' : [64, 'P', 128, 'P', 256, 256, 'P', 512, 512, 'P', 512, 512, 'P'], 16 | 'B' : [64, 64, 'P', 128, 128, 'P', 256, 256, 'P', 512, 512, 'P', 512, 512, 'P'], 17 | 'D' : [64, 64, 'P', 128, 128, 'P', 256, 256, 256, 'P', 512, 512, 512, 'P', 512, 512, 512, 'P'], 18 | 'E' : [64, 64, 'P', 128, 128, 'P', 256, 256, 256, 256, 'P', 512, 512, 512, 512, 'P', 512, 512, 512, 512, 'P'] 19 | } 20 | 21 | class VGG(nn.Module): 22 | 23 | def __init__(self, features, num_class=100): 24 | super().__init__() 25 | self.features = features 26 | 27 | self.classifier = nn.Sequential( 28 | nn.Linear(512, 4096), 29 | nn.ReLU(inplace=True), 30 | nn.Dropout(), 31 | nn.Linear(4096, 4096), 32 | nn.ReLU(inplace=True), 33 | nn.Dropout(), 34 | nn.Linear(4096, num_class) 35 | ) 36 | 37 | def forward(self, x): 38 | output = self.features(x) 39 | output = output.view(output.size()[0], -1) 40 | output = self.classifier(output) 41 | 42 | return output 43 | 44 | def make_layers(cfg, __C, batch_norm=False): 45 | layers = [] 46 | 47 | input_channel = 3 48 | for i,l in enumerate(cfg): 49 | if l == 'P': 50 | 51 | if __C.pooling == 'lip': 52 | layers += [get_pooling(__C)(cfg[i-1])] 53 | 54 | else: layers += [get_pooling(__C)(kernel_size=2,stride=2)] 55 | continue 56 | 57 | layers += [nn.Conv2d(input_channel, l, kernel_size=3, padding=1)] 58 | 59 | if batch_norm: 60 | layers += [nn.BatchNorm2d(l)] 61 | 62 | layers += [nn.ReLU(inplace=True)] 63 | input_channel = l 64 | 65 | return nn.Sequential(*layers) 66 | 67 | def vgg11_bn(__C): 68 | return VGG(make_layers(cfg['A'], __C, batch_norm=True)) 69 | 70 | def vgg13_bn(__C): 71 | return VGG(make_layers(cfg['B'], __C, batch_norm=True)) 72 | 73 | def vgg16_bn(__C): 74 | return VGG(make_layers(cfg['D'], __C, batch_norm=True)) 75 | 76 | def vgg19_bn(__C): 77 | return VGG(make_layers(cfg['E'], __C, batch_norm=True)) 78 | 79 | def vgg11(__C): 80 | return VGG(make_layers(cfg['A'], __C, batch_norm=False)) 81 | 82 | def vgg13(__C): 83 | return VGG(make_layers(cfg['B'], __C, batch_norm=False)) 84 | 85 | def vgg16(__C): 86 | return VGG(make_layers(cfg['D'], __C, batch_norm=False)) 87 | 88 | def vgg19(__C): 89 | return VGG(make_layers(cfg['E'], __C, batch_norm=False)) 90 | 91 | 92 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from configs import configs 3 | from utils.train_engine import train_engine 4 | 5 | def parse_args(): 6 | parser = argparse.ArgumentParser(description='Pooling Survey Args') 7 | parser.add_argument('--gpu', type=str, help="gpu choose, eg. '0,1,2,...' ") 8 | parser.add_argument('--run', type=str, dest='run_mode',choices=['train','test']) 9 | parser.add_argument('--name', type=str, required=True, help='the name of this training') 10 | parser.add_argument('--img_size', type=int, default=32, help='Resolution size') 11 | parser.add_argument('--batch_size', type=int, default=128, help='batch size for dataloader') 12 | parser.add_argument('--lr', type=float, default=0.1, help='initial learning rate') 13 | parser.add_argument('--lr_decay_rate', type=float, default=0.2, help="learning rate decay rate") 14 | parser.add_argument('--warmup_epoch', type=int, default=1, help='warmup epochs') 15 | parser.add_argument('--epoch', type=int, default=200, help='total epochs') 16 | parser.add_argument('--save_epoch', type=int, default=20, help="save model after every 20 epoch") 17 | parser.add_argument('--eval_every_epoch', action='store_true', default=True, help='evaluate the model every epoch') 18 | parser.add_argument('--pooling', type=str, 19 | choices=[ 20 | 'max', 21 | 'avg', 22 | 'mixed', 23 | 'Lp', 24 | 'lip', 25 | 'stochastic', 26 | 'soft' 27 | ], default='max',help='choose one pooling method to use', required=True) 28 | args = parser.parse_args() 29 | return args 30 | 31 | if __name__ == '__main__': 32 | args = parse_args() 33 | args_dict = configs.parse_to_dict(args) 34 | configs.add_args(args_dict) 35 | 36 | configs.training_init() 37 | configs.path_init() 38 | 39 | print("Hyper parameters:") 40 | print(configs) 41 | 42 | if configs.run_mode == 'train': 43 | train_engine(configs) 44 | -------------------------------------------------------------------------------- /utils/__pycache__/train_engine.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rentainhe/pytorch-pooling/f050e157adde95b2721136812158b014e993eee8/utils/__pycache__/train_engine.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/util.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rentainhe/pytorch-pooling/f050e157adde95b2721136812158b014e993eee8/utils/__pycache__/util.cpython-37.pyc -------------------------------------------------------------------------------- /utils/train_engine.py: -------------------------------------------------------------------------------- 1 | from models.net.vgg import vgg16_bn 2 | from models.net.vgg import vgg16 3 | from dataset_loader import get_train_loader 4 | from dataset_loader import get_test_loader 5 | import os 6 | import torch 7 | import torch.nn as nn 8 | import torch.optim as optim 9 | from utils.util import WarmUpLR 10 | from torch.utils.tensorboard import SummaryWriter 11 | import time 12 | def train_engine(__C): 13 | 14 | net = vgg16_bn(__C) 15 | net = net.cuda() 16 | 17 | # define dataloader 18 | train_loader = get_train_loader(__C) 19 | test_loader = get_test_loader(__C) 20 | 21 | # define optimizer and loss function 22 | loss_function = nn.CrossEntropyLoss() 23 | optimizer = optim.SGD(params=net.parameters(), lr=__C.lr, momentum=0.9, weight_decay=5e-4) 24 | 25 | # define optimizer scheduler 26 | train_scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=__C.milestones, gamma=__C.lr_decay_rate) 27 | iter_per_epoch = len(train_loader) 28 | warmup_schedule = WarmUpLR(optimizer, iter_per_epoch * __C.warmup_epoch) 29 | 30 | # define tensorboard writer 31 | writer = SummaryWriter(log_dir=os.path.join(__C.tensorboard_log_dir,__C.name,__C.version)) 32 | 33 | # define model save dir 34 | checkpoint_path = os.path.join(__C.ckpts_dir, __C.name, __C.version) 35 | if not os.path.exists(checkpoint_path): 36 | os.makedirs(checkpoint_path) 37 | checkpoint_path = os.path.join(checkpoint_path, '{net}-{epoch}-{type}.pth') 38 | 39 | # define the log save dir 40 | log_path = os.path.join(__C.result_log_dir, __C.name) 41 | if not os.path.exists(log_path): 42 | os.makedirs(log_path) 43 | log_path = os.path.join(log_path, __C.version + '.txt') 44 | 45 | # write the hyper parameters to log 46 | logfile = open(log_path, 'a+') 47 | logfile.write(str(__C)) 48 | logfile.close() 49 | 50 | best_acc = 0.0 51 | loss_sum = 0 52 | for epoch in range(1, __C.epoch): 53 | if epoch > __C.warmup_epoch: 54 | train_scheduler.step(epoch) 55 | 56 | start = time.time() 57 | net.train() 58 | for step, (images, labels) in enumerate(train_loader): 59 | if epoch <= __C.warmup_epoch: 60 | warmup_schedule.step() 61 | images = images.cuda() 62 | labels = labels.cuda() 63 | # using gradient accumulation step 64 | 65 | optimizer.zero_grad() 66 | loss_tmp = 0 67 | for accu_step in range(__C.gradient_accumulation_steps): 68 | loss_tmp = 0 69 | sub_images = images[accu_step * __C.sub_batch_size: 70 | (accu_step + 1) * __C.sub_batch_size] 71 | sub_labels = labels[accu_step * __C.sub_batch_size: 72 | (accu_step + 1) * __C.sub_batch_size] 73 | outputs = net(sub_images) 74 | loss = loss_function(outputs, sub_labels) 75 | loss.backward() 76 | # loss_tmp += loss.cpu().data.numpy() * __C.gradient_accumulation_steps 77 | # loss_sum += loss.cpu().data.numpy() * __C.gradient_accumulation_steps 78 | loss_tmp += loss.cpu().data.numpy() 79 | loss_sum += loss.cpu().data.numpy() 80 | 81 | optimizer.step() 82 | n_iter = (epoch-1) * len(train_loader) + step + 1 83 | print( 84 | '[{Version}] [{Name}] Training Epoch: {epoch} [{trained_samples}/{total_samples}]\tLoss: {:0.4f}\tLR: {:0.6f}'.format( 85 | loss_tmp, 86 | optimizer.param_groups[0]['lr'], 87 | Version=__C.version, 88 | Name=__C.name, 89 | epoch=epoch, 90 | trained_samples=step * __C.batch_size + len(images), 91 | total_samples=len(train_loader.dataset) 92 | )) 93 | # update training loss for each iteration 94 | 95 | writer.add_scalar('[Epoch] Train/loss', loss_tmp, n_iter) 96 | if epoch <= __C.warmup_epoch: 97 | writer.add_scalar('[Epoch] Train/lr', warmup_schedule.get_lr()[0], epoch) 98 | else: 99 | writer.add_scalar('[Epoch] Train/lr', train_scheduler.get_lr()[0], epoch) 100 | 101 | # update the result logfile 102 | logfile = open(log_path, 'a+') 103 | logfile.write( 104 | 'Epoch: ' + str(epoch) + 105 | ', Train Average Loss: {:.4f}'.format(loss_sum/len(train_loader.dataset)) + 106 | ', Lr: {:.6f}'.format(optimizer.param_groups[0]['lr']) + 107 | ', ' 108 | ) 109 | logfile.close() 110 | finish = time.time() 111 | print('epoch {} training time consumed: {:.2f}s'.format(epoch, finish - start)) 112 | 113 | if __C.eval_every_epoch: 114 | start = time.time() 115 | net.eval() 116 | test_loss = 0.0 117 | correct = 0.0 118 | for (images, labels) in test_loader: 119 | images = images.cuda() 120 | labels = labels.cuda() 121 | eval_outputs = net(images) 122 | eval_loss = loss_function(eval_outputs, labels) 123 | test_loss += eval_loss.item() 124 | _, preds = eval_outputs.max(1) 125 | correct += preds.eq(labels).sum() 126 | finish = time.time() 127 | 128 | test_average_loss = test_loss / len(test_loader.dataset) # 测试平均 loss 129 | acc = correct.float() / len(test_loader.dataset) # 测试准确率 130 | 131 | # save model after every "save_epoch" epoches and model with the best acc 132 | if epoch > __C.milestones[1] and best_acc < acc: 133 | torch.save(net.state_dict(), checkpoint_path.format(net=__C.name, epoch=epoch, type='best')) 134 | best_acc = acc 135 | continue 136 | if not epoch % __C.save_epoch: 137 | torch.save(net.state_dict(), checkpoint_path.format(net=__C.name, epoch=epoch, type='regular')) 138 | 139 | # print the testing information 140 | print('Evaluating Network.....') 141 | print('Test set: Average loss: {:.4f}, Accuracy: {:.4f}, Time consumed:{:.2f}s'.format( 142 | test_average_loss, 143 | acc, 144 | finish - start 145 | )) 146 | print() 147 | 148 | # update the result logfile 149 | logfile = open(log_path, 'a+') 150 | logfile.write( 151 | 'Test Average loss: {:.4f}'.format(test_average_loss) + 152 | ', Accuracy: {:.4f}'.format(acc) + 153 | '\n' 154 | ) 155 | logfile.close() 156 | 157 | # update the tensorboard log file 158 | writer.add_scalar('[Epoch] Test/Average loss', test_average_loss, epoch) 159 | writer.add_scalar('[Epoch] Test/Accuracy', acc, epoch) 160 | 161 | -------------------------------------------------------------------------------- /utils/util.py: -------------------------------------------------------------------------------- 1 | """ helper function 2 | 3 | author baiyu 4 | """ 5 | import os 6 | 7 | import re 8 | import datetime 9 | import numpy 10 | import torch.nn as nn 11 | from torch.optim.lr_scheduler import _LRScheduler 12 | 13 | 14 | def split_weights(net): 15 | """split network weights into to categlories, 16 | one are weights in conv layer and linear layer, 17 | others are other learnable paramters(conv bias, 18 | bn weights, bn bias, linear bias) 19 | 20 | Args: 21 | net: network architecture 22 | 23 | Returns: 24 | a dictionary of params splite into to categlories 25 | """ 26 | 27 | decay = [] 28 | no_decay = [] 29 | 30 | for m in net.modules(): 31 | if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear): 32 | decay.append(m.weight) 33 | 34 | if m.bias is not None: 35 | no_decay.append(m.bias) 36 | 37 | else: 38 | if hasattr(m, 'weight'): 39 | no_decay.append(m.weight) 40 | if hasattr(m, 'bias'): 41 | no_decay.append(m.bias) 42 | 43 | assert len(list(net.parameters())) == len(decay) + len(no_decay) 44 | 45 | return [dict(params=decay), dict(params=no_decay, weight_decay=0)] 46 | 47 | def compute_mean_std(cifar100_dataset): 48 | """compute the mean and std of cifar100 dataset 49 | Args: 50 | cifar100_training_dataset or cifar100_test_dataset 51 | witch derived from class torch.utils.data 52 | 53 | Returns: 54 | a tuple contains mean, std value of entire dataset 55 | """ 56 | 57 | data_r = numpy.dstack([cifar100_dataset[i][1][:, :, 0] for i in range(len(cifar100_dataset))]) 58 | data_g = numpy.dstack([cifar100_dataset[i][1][:, :, 1] for i in range(len(cifar100_dataset))]) 59 | data_b = numpy.dstack([cifar100_dataset[i][1][:, :, 2] for i in range(len(cifar100_dataset))]) 60 | mean = numpy.mean(data_r), numpy.mean(data_g), numpy.mean(data_b) 61 | std = numpy.std(data_r), numpy.std(data_g), numpy.std(data_b) 62 | 63 | return mean, std 64 | 65 | class WarmUpLR(_LRScheduler): 66 | """warmup_training learning rate scheduler 67 | Args: 68 | optimizer: optimzier(e.g. SGD) 69 | total_iters: totoal_iters of warmup phase 70 | """ 71 | def __init__(self, optimizer, total_iters, last_epoch=-1): 72 | 73 | self.total_iters = total_iters 74 | super().__init__(optimizer, last_epoch) 75 | 76 | def get_lr(self): 77 | """we will use the first m batches, and set the learning 78 | rate to base_lr * m / total_iters 79 | """ 80 | return [base_lr * self.last_epoch / (self.total_iters + 1e-8) for base_lr in self.base_lrs] 81 | 82 | 83 | def most_recent_folder(net_weights, fmt): 84 | """ 85 | return most recent created folder under net_weights 86 | if no none-empty folder were found, return empty folder 87 | """ 88 | # get subfolders in net_weights 89 | folders = os.listdir(net_weights) 90 | 91 | # filter out empty folders 92 | folders = [f for f in folders if len(os.listdir(os.path.join(net_weights, f)))] 93 | if len(folders) == 0: 94 | return '' 95 | 96 | # sort folders by folder created time 97 | folders = sorted(folders, key=lambda f: datetime.datetime.strptime(f, fmt)) 98 | return folders[-1] 99 | 100 | def most_recent_weights(weights_folder): 101 | """ 102 | return most recent created weights file 103 | if folder is empty return empty string 104 | """ 105 | weight_files = os.listdir(weights_folder) 106 | if len(weights_folder) == 0: 107 | return '' 108 | 109 | regex_str = r'([A-Za-z0-9]+)-([0-9]+)-(regular|best)' 110 | 111 | # sort files by epoch 112 | weight_files = sorted(weight_files, key=lambda w: int(re.search(regex_str, w).groups()[1])) 113 | 114 | return weight_files[-1] 115 | 116 | def last_epoch(weights_folder): 117 | weight_file = most_recent_weights(weights_folder) 118 | if not weight_file: 119 | raise Exception('no recent weights were found') 120 | resume_epoch = int(weight_file.split('-')[1]) 121 | 122 | return resume_epoch 123 | 124 | def best_acc_weights(weights_folder): 125 | """ 126 | return the best acc .pth file in given folder, if no 127 | best acc weights file were found, return empty string 128 | """ 129 | files = os.listdir(weights_folder) 130 | if len(files) == 0: 131 | return '' 132 | 133 | regex_str = r'([A-Za-z0-9]+)-([0-9]+)-(regular|best)' 134 | best_files = [w for w in files if re.search(regex_str, w).groups()[2] == 'best'] 135 | if len(best_files) == 0: 136 | return '' 137 | 138 | best_files = sorted(best_files, key=lambda w: int(re.search(regex_str, w).groups()[1])) 139 | return best_files[-1] --------------------------------------------------------------------------------