├── ClsName2id.txt ├── Nadam.py ├── README.md ├── config.py ├── data.py ├── data └── README.md ├── dataset.py ├── model.py ├── preprocess.py ├── test.py ├── train.py └── utils.py /ClsName2id.txt: -------------------------------------------------------------------------------- 1 | 旱地:dry-field:1 2 | 水田:paddy-field:2 3 | 梯田:terraced-field:3 4 | 草地:meadow:4 5 | 林地:forest:5 6 | 商业区:commercial-area:6 7 | 油田:oil-field:7 8 | 油罐区:storage-tank:8 9 | 工厂:works:9 10 | 矿区:mine:10 11 | 太阳能发电厂:solar-power-plant:11 12 | 风力发电站:wind-turbine:12 13 | 公园:park:13 14 | 游泳池:swimming-pool:14 15 | 教堂:church:15 16 | 墓地:cemetery:16 17 | 棒球场:baseball-field:17 18 | 篮球场:basketball-court:18 19 | 高尔夫球场:golf-course:19 20 | 足球场:soccer-field:20 21 | 温室:greenhouse:21 22 | 网球场:tennis-court:22 23 | 居民区:residential-area:23 24 | 岛屿:island:24 25 | 河流:river:25 26 | 停机坪:apron:26 27 | 直升机场:helipad:27 28 | 机场跑道:runway:28 29 | 桥梁:bridge:29 30 | 停车场:parking-lot:30 31 | 公路:road:31 32 | 路边停车区:roadside-parking-lot:32 33 | 转盘:roundabout:33 34 | 立交桥:viaduct:34 35 | 港口:port:35 36 | 铁路:railway:36 37 | 火车站:train-station:37 38 | 裸地:bare-land:38 39 | 沙漠:desert:39 40 | 冰岛:ice-land:40 41 | 山地:mountain:41 42 | 石质地:rock-land:42 43 | 稀疏灌木地:sparse-shrub-land:43 44 | 海滩:beach:44 45 | 湖泊:lake:45 46 | -------------------------------------------------------------------------------- /Nadam.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.optim import Optimizer 3 | import math 4 | 5 | class Nadam(Optimizer): 6 | 7 | def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, 8 | schedule_decay=0.004,amsgrad=False): 9 | if not 0.0 <= betas[0] < 1.0: 10 | raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0])) 11 | if not 0.0 <= betas[1] < 1.0: 12 | raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1])) 13 | defaults = dict(lr=lr, betas=betas, eps=eps, 14 | amsgrad=amsgrad,schedule_decay=schedule_decay) 15 | super(Nadam, self).__init__(params, defaults) 16 | 17 | def step(self, closure=None): 18 | loss = None 19 | if closure is not None: 20 | loss = closure() 21 | 22 | for group in self.param_groups: 23 | for p in group['params']: 24 | if p.grad is None: 25 | continue 26 | grad = p.grad.data 27 | if grad.is_sparse: 28 | raise RuntimeError('Nadam does not support sparse gradients, please consider SparseAdam instead') 29 | amsgrad = group['amsgrad'] 30 | 31 | state = self.state[p] 32 | 33 | # State initialization 34 | if len(state) == 0: 35 | state['step'] = 0 36 | # Exponential moving average of gradient values 37 | state['exp_avg'] = torch.zeros_like(p.data) 38 | # Exponential moving average of squared gradient values 39 | state['exp_avg_sq'] = torch.zeros_like(p.data) 40 | 41 | state['m_schedule'] = 1 42 | if amsgrad: 43 | # Maintains max of all exp. moving avg. of sq. grad. values 44 | state['max_exp_avg_sq'] = torch.zeros_like(p.data) 45 | 46 | exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] 47 | if amsgrad: 48 | max_exp_avg_sq = state['max_exp_avg_sq'] 49 | beta1, beta2 = group['betas'] 50 | 51 | 52 | state['step'] += 1 53 | momentum_cache_t = beta1 * ( 54 | 1. - 0.5 * math.pow(0.96, state['step'] * group['schedule_decay'] )) 55 | momentum_cache_t_1 = beta1 * ( 56 | 1. - 0.5 * math.pow(0.96, (state['step']+1) * group['schedule_decay'] )) 57 | state['m_schedule'] = state['m_schedule'] * momentum_cache_t 58 | 59 | exp_avg.mul_(beta1).add_(1 - beta1, grad) 60 | m_t_prime = exp_avg/(1 - state['m_schedule'] * momentum_cache_t_1) 61 | 62 | g_prime = grad.div(1 - state['m_schedule']) 63 | m_t_bar = (1. - momentum_cache_t) * g_prime + momentum_cache_t_1 * m_t_prime 64 | 65 | exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) 66 | if amsgrad: 67 | # Maintains the maximum of all 2nd moment running avg. till now 68 | torch.max(max_exp_avg_sq, exp_avg_sq , out=max_exp_avg_sq) 69 | # Use the max. for normalizing running avg. of gradient 70 | v_t_prime = max_exp_avg_sq/(1 - beta2 ** state['step']) 71 | else: 72 | v_t_prime = exp_avg_sq / (1 - beta2 ** state['step']) 73 | 74 | denom = v_t_prime.sqrt().add_(group['eps']) 75 | p.data.addcdiv_(-group['lr'], m_t_bar , denom) 76 | 77 | return loss -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # rssrai2019_scene_classification 2 | 3 | * The contest address is [here](http://rscup.bjxintong.com.cn/#/theme/1) 4 | 5 | * Baseline: Test Acc 90.14% 6 | 7 | ======= 8 | - step1 9 | 10 | ```python 11 | python3 preprocess.py 12 | ``` 13 | 14 | - step2 15 | 16 | ​ Train from scratch: 17 | 18 | ```python 19 | python3 train.py 20 | ``` 21 | 22 | ​ or Train from checkpoint: 23 | 24 | ```python 25 | python3 train.py --net-params /path/to/*.pth 26 | ``` 27 | 28 | - step3 29 | 30 | ```python 31 | python3 test.py 32 | ``` 33 | 34 | -------------------------------------------------------------------------------- /config.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author : Magic 3 | # @Time : 2019/7/4 12:01 4 | # @File : config.py 5 | 6 | config_dict = dict() 7 | 8 | config_dict['name_to_id'] = 'ClsName2id.txt' 9 | config_dict['data_dir'] = 'data' 10 | config_dict['data_dir_train'] = 'data/train' 11 | config_dict['data_dir_val'] = 'data/val' 12 | config_dict['data_dir_test'] = 'data/test' 13 | config_dict['im_size'] = (224, 224) 14 | 15 | 16 | -------------------------------------------------------------------------------- /data.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author : Magic 3 | # @Time : 2019/7/4 12:01 4 | # @File : data.py 5 | 6 | from config import config_dict 7 | from dataset import SenseData, SenseDataTest 8 | from torchvision import transforms 9 | 10 | img_train = config_dict['data_dir_train'] 11 | img_val = config_dict['data_dir_val'] 12 | img_test = config_dict['data_dir_test'] 13 | 14 | def train_augs(): 15 | return transforms.Compose([ 16 | transforms.RandomResizedCrop(config_dict['im_size']), 17 | transforms.RandomHorizontalFlip(), 18 | transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1), 19 | transforms.ToTensor(), 20 | transforms.Normalize( 21 | [0.485, 0.456, 0.406], 22 | [0.229, 0.224, 0.225] 23 | ) 24 | ]) 25 | 26 | def val_augs(): 27 | return transforms.Compose([ 28 | transforms.Resize(256, interpolation=2), 29 | transforms.CenterCrop(config_dict['im_size']), 30 | transforms.ToTensor(), 31 | transforms.Normalize( 32 | [0.485, 0.456, 0.406], 33 | [0.229, 0.224, 0.225] 34 | ) 35 | ]) 36 | 37 | def test_augs(): 38 | return transforms.Compose([ 39 | transforms.Resize(256, interpolation=2), 40 | transforms.CenterCrop(config_dict['im_size']), 41 | transforms.ToTensor(), 42 | transforms.Normalize( 43 | [0.485, 0.456, 0.406], 44 | [0.229, 0.224, 0.225] 45 | ) 46 | ]) 47 | def get_train_data(img_path=img_train, transform = train_augs()): 48 | return SenseData(img_path, transform) 49 | 50 | def get_val_data(img_path=img_val, transform = val_augs()): 51 | return SenseData(img_path, transform) 52 | 53 | def get_test_data(img_path=img_test, transform = test_augs()): 54 | return SenseDataTest(img_path, transform) 55 | 56 | -------------------------------------------------------------------------------- /data/README.md: -------------------------------------------------------------------------------- 1 | put data in here 2 | 3 | Like this. 4 | 5 | ├── data 6 | 7 | │ ├── train 8 | 9 | │ ├── val 10 | 11 | │ ├── test 12 | 13 | 14 | -------------------------------------------------------------------------------- /dataset.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author : Magic 3 | # @Time : 2019/7/5 10:23 4 | # @File : dataset.py 5 | 6 | import torch 7 | import glob 8 | from torch.utils.data import Dataset,DataLoader 9 | from torchvision import transforms 10 | from config import config_dict 11 | from utils import map_label 12 | from PIL import Image 13 | import cv2 14 | 15 | class SenseData(Dataset): 16 | def __init__(self, img_dir, transform = None): 17 | super(SenseData, self).__init__() 18 | self.img_list = glob.glob(img_dir+'/*/*.jpg') 19 | self.transform = transform 20 | 21 | def __len__(self): 22 | return len(self.img_list) 23 | 24 | def __getitem__(self, index): 25 | img = self.img_list[index] 26 | _, label_map = map_label(config_dict['name_to_id']) 27 | img_label = label_map[img.split('/')[2]] 28 | img = Image.open(img) 29 | if self.transform: 30 | img = self.transform(img) 31 | 32 | return img, torch.Tensor([int(img_label)-1]) 33 | 34 | 35 | class SenseDataTest(Dataset): 36 | def __init__(self, img_dir, transform = None): 37 | super(SenseDataTest, self).__init__() 38 | self.img_list = sorted(glob.glob(img_dir+'/*.jpg')) 39 | self.transform = transform 40 | 41 | def __len__(self): 42 | return len(self.img_list) 43 | 44 | def __getitem__(self, index): 45 | img = self.img_list[index] 46 | filedir=img 47 | img = Image.open(img) 48 | if self.transform: 49 | img = self.transform(img) 50 | return index, filedir, img 51 | 52 | if __name__ == '__main__': 53 | # train_data = SenseData(config_dict['data_dir_train'], None) 54 | # train_loader = DataLoader(train_data, batch_size=4, shuffle=False, num_workers=1) 55 | # for i, batch in enumerate(train_loader): 56 | # print(batch[0], batch[1]) 57 | test_data = SenseDataTest(config_dict['data_dir_test'], transforms.Compose([transforms.Resize(224, interpolation=2),transforms.ToTensor()])) 58 | test_loader = DataLoader(test_data, batch_size=1, shuffle=False, num_workers=1) 59 | for i, batch in enumerate(test_loader): 60 | print(batch.shape) 61 | break -------------------------------------------------------------------------------- /model.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author : Magic 3 | # @Time : 2019/7/4 12:02 4 | # @File : model.py 5 | 6 | """ 7 | SENet, implemented in PyTorch. 8 | Original paper: 'Squeeze-and-Excitation Networks,' https://arxiv.org/abs/1709.01507. 9 | """ 10 | 11 | __all__ = ['SENet', 'senet52', 'senet103', 'senet154', 'SEInitBlock'] 12 | 13 | import os 14 | import math 15 | import torch.nn as nn 16 | import torch.nn.init as init 17 | from common import conv1x1_block, conv3x3_block, SEBlock 18 | 19 | 20 | class SENetBottleneck(nn.Module): 21 | """ 22 | SENet bottleneck block for residual path in SENet unit. 23 | 24 | Parameters: 25 | ---------- 26 | in_channels : int 27 | Number of input channels. 28 | out_channels : int 29 | Number of output channels. 30 | stride : int or tuple/list of 2 int 31 | Strides of the convolution. 32 | cardinality: int 33 | Number of groups. 34 | bottleneck_width: int 35 | Width of bottleneck block. 36 | """ 37 | def __init__(self, 38 | in_channels, 39 | out_channels, 40 | stride, 41 | cardinality, 42 | bottleneck_width): 43 | super(SENetBottleneck, self).__init__() 44 | mid_channels = out_channels // 4 45 | D = int(math.floor(mid_channels * (bottleneck_width / 64.0))) 46 | group_width = cardinality * D 47 | group_width2 = group_width // 2 48 | 49 | self.conv1 = conv1x1_block( 50 | in_channels=in_channels, 51 | out_channels=group_width2) 52 | self.conv2 = conv3x3_block( 53 | in_channels=group_width2, 54 | out_channels=group_width, 55 | stride=stride, 56 | groups=cardinality) 57 | self.conv3 = conv1x1_block( 58 | in_channels=group_width, 59 | out_channels=out_channels, 60 | activate=False) 61 | 62 | def forward(self, x): 63 | x = self.conv1(x) 64 | x = self.conv2(x) 65 | x = self.conv3(x) 66 | return x 67 | 68 | 69 | class SENetUnit(nn.Module): 70 | """ 71 | SENet unit. 72 | 73 | Parameters: 74 | ---------- 75 | in_channels : int 76 | Number of input channels. 77 | out_channels : int 78 | Number of output channels. 79 | stride : int or tuple/list of 2 int 80 | Strides of the convolution. 81 | cardinality: int 82 | Number of groups. 83 | bottleneck_width: int 84 | Width of bottleneck block. 85 | identity_conv3x3 : bool, default False 86 | Whether to use 3x3 convolution in the identity link. 87 | """ 88 | def __init__(self, 89 | in_channels, 90 | out_channels, 91 | stride, 92 | cardinality, 93 | bottleneck_width, 94 | identity_conv3x3): 95 | super(SENetUnit, self).__init__() 96 | self.resize_identity = (in_channels != out_channels) or (stride != 1) 97 | 98 | self.body = SENetBottleneck( 99 | in_channels=in_channels, 100 | out_channels=out_channels, 101 | stride=stride, 102 | cardinality=cardinality, 103 | bottleneck_width=bottleneck_width) 104 | self.se = SEBlock(channels=out_channels) 105 | if self.resize_identity: 106 | if identity_conv3x3: 107 | self.identity_conv = conv3x3_block( 108 | in_channels=in_channels, 109 | out_channels=out_channels, 110 | stride=stride, 111 | activate=False) 112 | else: 113 | self.identity_conv = conv1x1_block( 114 | in_channels=in_channels, 115 | out_channels=out_channels, 116 | stride=stride, 117 | activate=False) 118 | self.activ = nn.ReLU(inplace=True) 119 | 120 | def forward(self, x): 121 | if self.resize_identity: 122 | identity = self.identity_conv(x) 123 | else: 124 | identity = x 125 | x = self.body(x) 126 | x = self.se(x) 127 | x = x + identity 128 | x = self.activ(x) 129 | return x 130 | 131 | 132 | class SEInitBlock(nn.Module): 133 | """ 134 | SENet specific initial block. 135 | 136 | Parameters: 137 | ---------- 138 | in_channels : int 139 | Number of input channels. 140 | out_channels : int 141 | Number of output channels. 142 | """ 143 | def __init__(self, 144 | in_channels, 145 | out_channels): 146 | super(SEInitBlock, self).__init__() 147 | mid_channels = out_channels // 2 148 | 149 | self.conv1 = conv3x3_block( 150 | in_channels=in_channels, 151 | out_channels=mid_channels, 152 | stride=2) 153 | self.conv2 = conv3x3_block( 154 | in_channels=mid_channels, 155 | out_channels=mid_channels) 156 | self.conv3 = conv3x3_block( 157 | in_channels=mid_channels, 158 | out_channels=out_channels) 159 | self.pool = nn.MaxPool2d( 160 | kernel_size=3, 161 | stride=2, 162 | padding=1) 163 | 164 | def forward(self, x): 165 | x = self.conv1(x) 166 | x = self.conv2(x) 167 | x = self.conv3(x) 168 | x = self.pool(x) 169 | return x 170 | 171 | 172 | class SENet(nn.Module): 173 | """ 174 | SENet model from 'Squeeze-and-Excitation Networks,' https://arxiv.org/abs/1709.01507. 175 | 176 | Parameters: 177 | ---------- 178 | channels : list of list of int 179 | Number of output channels for each unit. 180 | init_block_channels : int 181 | Number of output channels for the initial unit. 182 | cardinality: int 183 | Number of groups. 184 | bottleneck_width: int 185 | Width of bottleneck block. 186 | in_channels : int, default 3 187 | Number of input channels. 188 | in_size : tuple of two ints, default (224, 224) 189 | Spatial size of the expected input image. 190 | num_classes : int, default 1000 191 | Number of classification classes. 192 | """ 193 | def __init__(self, 194 | channels, 195 | init_block_channels, 196 | cardinality, 197 | bottleneck_width, 198 | in_channels=3, 199 | in_size=(224, 224), 200 | num_classes=1000): 201 | super(SENet, self).__init__() 202 | self.in_size = in_size 203 | self.num_classes = num_classes 204 | 205 | self.features = nn.Sequential() 206 | self.features.add_module("init_block", SEInitBlock( 207 | in_channels=in_channels, 208 | out_channels=init_block_channels)) 209 | in_channels = init_block_channels 210 | for i, channels_per_stage in enumerate(channels): 211 | stage = nn.Sequential() 212 | identity_conv3x3 = (i != 0) 213 | for j, out_channels in enumerate(channels_per_stage): 214 | stride = 2 if (j == 0) and (i != 0) else 1 215 | stage.add_module("unit{}".format(j + 1), SENetUnit( 216 | in_channels=in_channels, 217 | out_channels=out_channels, 218 | stride=stride, 219 | cardinality=cardinality, 220 | bottleneck_width=bottleneck_width, 221 | identity_conv3x3=identity_conv3x3)) 222 | in_channels = out_channels 223 | self.features.add_module("stage{}".format(i + 1), stage) 224 | self.features.add_module("final_pool", nn.AvgPool2d( 225 | kernel_size=7, 226 | stride=1)) 227 | 228 | self.output = nn.Sequential() 229 | self.output.add_module("dropout", nn.Dropout(p=0.2)) 230 | self.output.add_module("fc", nn.Linear( 231 | in_features=in_channels, 232 | out_features=num_classes)) 233 | 234 | self._init_params() 235 | 236 | def _init_params(self): 237 | for name, module in self.named_modules(): 238 | if isinstance(module, nn.Conv2d): 239 | init.kaiming_uniform_(module.weight) 240 | if module.bias is not None: 241 | init.constant_(module.bias, 0) 242 | 243 | def forward(self, x): 244 | x = self.features(x) 245 | x = x.view(x.size(0), -1) 246 | x = self.output(x) 247 | return x 248 | 249 | 250 | def get_senet(blocks, 251 | model_name=None, 252 | pretrained=False, 253 | root=os.path.join("~", ".torch", "models"), 254 | **kwargs): 255 | """ 256 | Create SENet model with specific parameters. 257 | 258 | Parameters: 259 | ---------- 260 | blocks : int 261 | Number of blocks. 262 | model_name : str or None, default None 263 | Model name for loading pretrained model. 264 | pretrained : bool, default False 265 | Whether to load the pretrained weights for model. 266 | root : str, default '~/.torch/models' 267 | Location for keeping the model parameters. 268 | """ 269 | 270 | if blocks == 52: 271 | layers = [3, 4, 6, 3] 272 | cardinality = 32 273 | elif blocks == 103: 274 | layers = [3, 4, 23, 3] 275 | cardinality = 32 276 | elif blocks == 154: 277 | layers = [3, 8, 36, 3] 278 | cardinality = 64 279 | else: 280 | raise ValueError("Unsupported SENet with number of blocks: {}".format(blocks)) 281 | 282 | bottleneck_width = 4 283 | init_block_channels = 128 284 | channels_per_layers = [256, 512, 1024, 2048] 285 | 286 | channels = [[ci] * li for (ci, li) in zip(channels_per_layers, layers)] 287 | 288 | net = SENet( 289 | channels=channels, 290 | init_block_channels=init_block_channels, 291 | cardinality=cardinality, 292 | bottleneck_width=bottleneck_width, 293 | **kwargs) 294 | 295 | if pretrained: 296 | if (model_name is None) or (not model_name): 297 | raise ValueError("Parameter `model_name` should be properly initialized for loading pretrained model.") 298 | from .model_store import download_model 299 | download_model( 300 | net=net, 301 | model_name=model_name, 302 | local_model_store_dir_path=root) 303 | 304 | return net 305 | 306 | 307 | def senet52(**kwargs): 308 | """ 309 | SENet-52 model from 'Squeeze-and-Excitation Networks,' https://arxiv.org/abs/1709.01507. 310 | 311 | Parameters: 312 | ---------- 313 | pretrained : bool, default False 314 | Whether to load the pretrained weights for model. 315 | root : str, default '~/.torch/models' 316 | Location for keeping the model parameters. 317 | """ 318 | return get_senet(blocks=52, model_name="senet52", **kwargs) 319 | 320 | 321 | def senet103(**kwargs): 322 | """ 323 | SENet-103 model from 'Squeeze-and-Excitation Networks,' https://arxiv.org/abs/1709.01507. 324 | 325 | Parameters: 326 | ---------- 327 | pretrained : bool, default False 328 | Whether to load the pretrained weights for model. 329 | root : str, default '~/.torch/models' 330 | Location for keeping the model parameters. 331 | """ 332 | return get_senet(blocks=103, model_name="senet103", **kwargs) 333 | 334 | 335 | def senet154(**kwargs): 336 | """ 337 | SENet-154 model from 'Squeeze-and-Excitation Networks,' https://arxiv.org/abs/1709.01507. 338 | 339 | Parameters: 340 | ---------- 341 | pretrained : bool, default False 342 | Whether to load the pretrained weights for model. 343 | root : str, default '~/.torch/models' 344 | Location for keeping the model parameters. 345 | """ 346 | return get_senet(blocks=154, model_name="senet154", **kwargs) 347 | 348 | 349 | def _calc_width(net): 350 | import numpy as np 351 | net_params = filter(lambda p: p.requires_grad, net.parameters()) 352 | weight_count = 0 353 | for param in net_params: 354 | weight_count += np.prod(param.size()) 355 | return weight_count 356 | 357 | 358 | def _test(): 359 | import torch 360 | from torch.autograd import Variable 361 | 362 | pretrained = False 363 | 364 | models = [ 365 | senet52, 366 | senet103, 367 | senet154, 368 | ] 369 | 370 | for model in models: 371 | 372 | net = model(pretrained=pretrained) 373 | 374 | # net.train() 375 | net.eval() 376 | weight_count = _calc_width(net) 377 | print("m={}, {}".format(model.__name__, weight_count)) 378 | assert (model != senet52 or weight_count == 44659416) 379 | assert (model != senet103 or weight_count == 60963096) 380 | assert (model != senet154 or weight_count == 115088984) 381 | 382 | x = Variable(torch.randn(1, 3, 224, 224)) 383 | y = net(x) 384 | y.sum().backward() 385 | assert (tuple(y.size()) == (1, 1000)) 386 | 387 | 388 | if __name__ == "__main__": 389 | _test() 390 | 391 | 392 | -------------------------------------------------------------------------------- /preprocess.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author : Magic 3 | # @Time : 2019/7/4 12:04 4 | # @File : preprocess.py 5 | 6 | import os 7 | import utils 8 | from config import config_dict 9 | 10 | # transfrom cn_name to en_name 11 | def transform_name(path): 12 | cn_to_en, _ = utils.map_label() 13 | for ro, di, fi in os.walk(path): 14 | dirname = os.path.dirname(ro) 15 | name = ro.split('/')[-1] 16 | print(dirname, name) 17 | if name in cn_to_en: 18 | os.rename(os.path.join(dirname, name), os.path.join(dirname, cn_to_en[name])) 19 | print('Rename Success!') 20 | 21 | def calc_trainset_mean_std(): 22 | pass 23 | 24 | if __name__ == '__main__': 25 | transform_name(config_dict['data_dir']) -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author : Magic 3 | # @Time : 2019/7/4 12:02 4 | # @File : test.py 5 | 6 | import torch 7 | import argparse, os 8 | import numpy as np 9 | import torchvision.models as models 10 | from torch.utils.data import DataLoader 11 | from torch.autograd import Variable 12 | from data import get_test_data, get_val_data 13 | 14 | parser = argparse.ArgumentParser(description='test scene data...') 15 | parser.add_argument('--gpus', default='0', type=str, help='ordinates of gpus to use, can be "0,1,2,3" ') 16 | parser.add_argument('--batch-size', default=1, type=int, help='batch size') 17 | parser.add_argument('--num-workers', default=4, type=int, help='the num of threads to load data') 18 | parser.add_argument('--resume', type=str, default='senet_model_29.pth') 19 | 20 | args = parser.parse_args() 21 | os.environ['CUDA_VISBLE_DEVICES'] = args.gpus 22 | 23 | # Load Datasets 24 | val_data = get_val_data() 25 | test_data = get_test_data() 26 | 27 | val_loader = DataLoader(val_data, num_workers=args.num_workers, batch_size=128, shuffle=False) 28 | test_loader = DataLoader(test_data, num_workers=args.num_workers, batch_size=1, shuffle=False) 29 | 30 | # Build Model 31 | model = models.resnext50_32x4d() 32 | in_feature = model.fc.in_features 33 | model.fc = torch.nn.Linear(in_feature, 45) 34 | model_dict = torch.load(args.resume)['model'] 35 | pred_dict = {} 36 | for k,v in model_dict.items(): 37 | pred_dict[k.replace('module.','')] = v 38 | 39 | #model_dict.update(pred_dict) 40 | model.load_state_dict(pred_dict) 41 | model.to('cuda:0') 42 | model.eval() 43 | res = [] 44 | 45 | 46 | def val(): 47 | with torch.no_grad(): 48 | rightN = 0 49 | for idx, batch in enumerate(val_loader, 1): 50 | img, label = Variable(batch[0]), Variable(torch.from_numpy(np.array(batch[1])).long()) 51 | #print(img.shape, label.shape) 52 | if torch.cuda.is_available(): 53 | img = img.to('cuda:0') 54 | label = label.to('cuda:0') 55 | pred = model(img) 56 | 57 | # cal acc 58 | pred = np.argmax(pred.data.cpu().numpy(), axis=1) 59 | gt = label.squeeze().cpu().numpy() 60 | rightN += (pred == gt).sum() 61 | print('==> Complete. Acc: {:.4f} '.format(rightN / len(val_loader.dataset))) 62 | 63 | 64 | def eval(): 65 | with torch.no_grad(): 66 | for idx, batch in enumerate(test_loader): 67 | #print(batch[0], batch[1]) 68 | pred = model(batch[2]) 69 | pred_label = pred.argmax(dim=1) 70 | result = '{:05d}.jpg {}'.format(idx+1, int(pred_label.item()) + 1) 71 | print(result + ' is done!') 72 | 73 | #print('\n') 74 | res.append(result) 75 | 76 | with open('submit.txt', 'w') as f: 77 | for line in res: 78 | f.write(line + '\n') 79 | 80 | 81 | if __name__ == '__main__': 82 | val() 83 | #eval() 84 | # print('Submit.txt is Finished!') 85 | 86 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author : Magic 3 | # @Time : 2019/7/4 12:02 4 | # @File : train.py 5 | 6 | import torch, time 7 | import torchvision.models as models 8 | import argparse 9 | import os 10 | import numpy as np 11 | from torch.utils.data import DataLoader 12 | import torch.nn as nn 13 | import torch.optim as optim 14 | from torch.autograd import Variable 15 | from tensorboardX import SummaryWriter 16 | from Nadam import Nadam 17 | from data import get_train_data, get_val_data 18 | from utils import save_checkpoint 19 | import utils 20 | 21 | 22 | def parse_args(): 23 | parser = argparse.ArgumentParser(description='train scene data') 24 | parser.add_argument('--batch-size', default=128, type=int, help='batch size') 25 | parser.add_argument('--num-workers', default=4, type=int ,help='the num of threads to load data') 26 | parser.add_argument('--epochs', default=100, type=int, help='training epochs, Default 500') 27 | parser.add_argument('--lr', default=0.01, type=float, help='learning rate') 28 | parser.add_argument('--num-classes', default=45, type=int, help='the number of classes') 29 | parser.add_argument('--gpus', default='0,1', type=str, help='ordinates of gpus to use, can be "0,1,2,3" ') 30 | parser.add_argument('--seed', default=666, type=int, help='random seed to use, Default=666') 31 | 32 | parser.add_argument('--begin-epoch', default=0, type=int, help='begin epoch') 33 | parser.add_argument('--lr-factor', default=0.1, type=float, help='the ratio to reduce lr on each step') 34 | parser.add_argument('--lr-step-epochs', default='20,45,60,80', type=str, help='the epochs to reduce the lr') 35 | parser.add_argument('--save-model-prefix', default='resnext', type=str, help='model prefix') 36 | parser.add_argument('--save-model-step', type=int, default=1, help='snapshot step (epoch num)') 37 | 38 | parser.add_argument('--net-params', type=str, default=None, help='resume the training') 39 | parser.add_argument('--log-dir', type=str, default='log', help='the directory of the log') 40 | parser.add_argument('--log-file', type=str, default='log.txt', help='log file path') 41 | 42 | return parser.parse_args() 43 | 44 | args = parse_args() 45 | writer = SummaryWriter(log_dir='logs_board/resnext') 46 | devs = [int(x) for x in args.gpus.split(',')] 47 | lr_step_epochs = [int(x) for x in args.lr_step_epochs.split(',')] 48 | if args.log_dir: 49 | utils.create_dir(args.log_dir) 50 | logger = utils.Logger(os.path.join(args.log_dir, args.log_file)) 51 | 52 | train_data = get_train_data() 53 | train_loader = DataLoader(train_data, batch_size=args.batch_size, num_workers=args.num_workers, shuffle=True) 54 | 55 | model = models.resnext50_32x4d(pretrained=True) 56 | num_fc = model.fc.in_features 57 | model.fc = nn.Linear(num_fc, args.num_classes) 58 | 59 | if args.net_params: 60 | print('=> Loading checkpoint... ') 61 | resume_model = torch.load(args.net_params) 62 | model_dict = resume_model['model'] 63 | args.begin_epoch = resume_model['epoch'] 64 | pred_dict = {} 65 | for k,v in model_dict.items(): 66 | pred_dict[k.replace('module.','')] = v 67 | model.load_state_dict(pred_dict) 68 | 69 | if args.begin_epoch: 70 | for i in lr_step_epochs: 71 | if args.begin_epoch>=i: 72 | args.lr = args.lr*0.1 73 | print('Learning rate is ', args.lr) 74 | 75 | model = nn.DataParallel(model, device_ids=devs) 76 | model.to('cuda:0') 77 | criterion = nn.CrossEntropyLoss() 78 | optimizer = Nadam(model.parameters(), lr=args.lr) 79 | 80 | val_data = get_val_data() 81 | val_loader = DataLoader(val_data, batch_size=args.batch_size, num_workers=args.num_workers, shuffle=False) 82 | 83 | def train(epoch): 84 | epoch_loss, rightN = 0, 0 85 | model.train() 86 | for idx, batch in enumerate(train_loader, 1): 87 | img, label = Variable(batch[0], requires_grad=True), Variable(torch.from_numpy(np.array(batch[1])).long()) 88 | if torch.cuda.is_available(): 89 | img = img.to('cuda:0') 90 | label = label.to('cuda:0') 91 | 92 | optimizer.zero_grad() 93 | t0 = time.time() 94 | pred = model(img) 95 | #print(pred.shape, label.shape, label.squeeze().shape) 96 | loss = criterion(pred, label.squeeze()) 97 | 98 | #cal acc 99 | pred = np.argmax(pred.data.cpu().numpy(), axis=1) 100 | gt = label.squeeze().cpu().numpy() 101 | rightN += (pred==gt).sum() 102 | 103 | epoch_loss += loss.item() 104 | loss.backward() 105 | optimizer.step() 106 | writer.add_scalar('scalar/loss',loss.item(), epoch*len(train_loader)+idx) 107 | msg = '==> Epoch[{}]({}/{}): Loss: {:.4f} || Timer: {:.4f} sec.'.format(epoch, idx, len(train_loader),\ 108 | loss.item(), time.time()-t0) 109 | logger.write(msg) 110 | msg = '==> Epoch {} Complete. Train Acc: {:.4f} || Avg Loss: {:.4f}'.format(epoch, rightN/len(train_loader.dataset), epoch_loss/len(train_loader)) 111 | logger.write(msg) 112 | writer.add_scalar('scalar/train_acc', rightN/len(train_loader.dataset), epoch) 113 | 114 | def val(epoch): 115 | model.eval() 116 | with torch.no_grad(): 117 | count = 0 118 | for idx, batch in enumerate(val_loader, 1): 119 | img, label = Variable(batch[0]), Variable(torch.from_numpy(np.array(batch[1])).long()) 120 | if torch.cuda.is_available(): 121 | img = img.to('cuda:0') 122 | label = label.to('cuda:0') 123 | pred = model(img) 124 | 125 | #cal acc 126 | pred = np.argmax(pred.data.cpu().numpy(), axis=1) 127 | gt = label.squeeze().cpu().numpy() 128 | count += (pred==gt).sum() 129 | msg = '==> Train{}: Complete. Val Acc: {:.4f} '.format(epoch, count/len(val_loader.dataset)) 130 | logger.write(msg) 131 | writer.add_scalar('scalar/val_acc', count/len(val_loader.dataset), epoch) 132 | 133 | if __name__ == '__main__': 134 | for epoch in range(args.begin_epoch, args.epochs + 1): 135 | train(epoch) 136 | val(epoch) 137 | if epoch in lr_step_epochs: 138 | for param_group in optimizer.param_groups: 139 | param_group['lr'] *= args.lr_factor 140 | print('Learning rate decay : lr = {}'.format(optimizer.param_groups[0]['lr'])) 141 | 142 | if (epoch+1) % args.save_model_step == 0: 143 | save_checkpoint(model, epoch, args.save_model_prefix) 144 | 145 | writer.close() 146 | -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author : Magic 3 | # @Time : 2019/7/4 10:49 4 | # @File : utils.py 5 | 6 | import os 7 | import torch 8 | import numpy as np 9 | from config import config_dict 10 | #file_path = 'F:\\ai_competition\\rssrai2019_scene_classification\\ClsName2id.txt' 11 | file_path = config_dict['name_to_id'] 12 | 13 | def map_label(file_path=file_path): 14 | chinese_to_english = {} 15 | label_map = {} 16 | with open(file_path, encoding='utf-8') as f: 17 | for line in f.readlines(): 18 | cn_name, en_name, label = line.strip().split(':') 19 | chinese_to_english[cn_name] = en_name 20 | label_map[en_name] = label 21 | return chinese_to_english, label_map 22 | 23 | def create_dir(path): 24 | if not os.path.exists(path): 25 | try: 26 | os.makedirs(path) 27 | except: 28 | print('Create dir failed! try again.') 29 | raise 30 | 31 | def cuda(x): 32 | if torch.cuda.is_available(): 33 | if isinstance(x, (list, tuple)): 34 | return [_x.cuda() for _x in x] 35 | else: 36 | return x.cuda() 37 | 38 | def save_checkpoint(model, epoch, prefix): 39 | output_path = 'checkpoint/' + prefix + '_model_{}.pth'.format(epoch) 40 | if not os.path.exists('checkpoint/'): 41 | os.mkdir('checkpoint/') 42 | state = {'epoch': epoch, 'model':model.state_dict()} 43 | torch.save(state, output_path) 44 | print('Checkpoint save to {}'.format(output_path)) 45 | 46 | class Logger(object): 47 | def __init__(self, output_name): 48 | dirname = os.path.dirname(output_name) 49 | if not os.path.exists(dirname): 50 | os.makedirs(dirname) 51 | self.log_file = open(output_name, 'w') 52 | self.info = {} 53 | 54 | def append(self, key, value): 55 | vals = self.info.setdefault(key, []) 56 | vals.append(value) 57 | 58 | def log(self, extra_msg=''): 59 | msgs = [extra_msg] 60 | for key, vals in self.info.items(): 61 | msgs.append('%s %.6f' % (key, np.mean(vals))) 62 | msg = '\n'.join(msgs) 63 | self.log_file.write(msg + '\n') 64 | self.log_file.flush() 65 | self.info = {} 66 | return msg 67 | 68 | def write(self, msg): 69 | self.log_file.write(msg + '\n') 70 | self.log_file.flush() 71 | print(msg) 72 | 73 | 74 | 75 | 76 | --------------------------------------------------------------------------------