├── README.md ├── config.py ├── dataset ├── __init__.py ├── __pycache__ │ ├── aug.cpython-36.pyc │ └── dataloader.cpython-36.pyc ├── aug.py └── dataloader.py ├── main.py ├── models ├── __init__.py └── model.py ├── test.py └── utils ├── __init__.py ├── progress_bar.py └── utils.py /README.md: -------------------------------------------------------------------------------- 1 | # pytorch-image-classification 2 | use pytorch to do image classfiication tasks 3 | 4 | ### 使用方法 5 | 6 | 移步：[从实例掌握 pytorch 进行图像分类](http://spytensor.com/index.php/archives/21/) 7 | 8 | ## 2019.04.28 9 | add test.py 10 | 11 | #### version0.2.0 12 | 13 | **更新** 14 | 手写了一份进度条工具，效果如下 15 | 16 | ``` 17 | loading train dataset 18 | 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4572/4572 [00:00<00:00, 1145746.42it/s] 19 | loading train dataset 20 | 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2520/2520 [00:00<00:00, 1128994.45it/s] 21 | Train Epoch: 1/40 [>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>] [Current: Loss 1.079473 Top1: 75.131233 ] 286/286 [ 100% ] 22 | Val Epoch: 1/40 [>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>] [Current: Loss 0.469368 Top1: 89.484131 ] 79/79 [ 100% ] 23 | Get Better top1 : 89.4841 saving weights to ./checkpoints/best_model/resnet50/0/model_best.pth.tar 24 | Train Epoch: 2/40 [>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>] [Current: Loss 0.696026 Top1: 83.442696 ] 286/286 [ 100% ] 25 | Val Epoch: 2/40 [>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>] [Current: Loss 0.275893 Top1: 91.269844 ] 79/79 [ 100% ] 26 | Get Better top1 : 91.2698 saving weights to ./checkpoints/best_model/resnet50/0/model_best.pth.tar 27 | Train Epoch: 3/40 [>>>>> ] [Current: Loss 0.667610 Top1: 84.165291 ] 31/286 [ 10% ] 28 | ``` 29 | -------------------------------------------------------------------------------- /config.py: -------------------------------------------------------------------------------- 1 | class DefaultConfigs(object): 2 | #1.string parameters 3 | train_data = "/home/user/zcj/tutorials/data/train/" 4 | test_data = "" 5 | val_data = "/home/user/zcj/tutorials/data/val/" 6 | model_name = "resnet50" 7 | weights = "./checkpoints/" 8 | best_models = weights + "best_model/" 9 | submit = "./submit/" 10 | logs = "./logs/" 11 | gpus = "1" 12 | 13 | #2.numeric parameters 14 | epochs = 40 15 | batch_size = 16 16 | img_height = 300 17 | img_weight = 300 18 | num_classes = 62 19 | seed = 888 20 | lr = 1e-4 21 | lr_decay = 1e-4 22 | weight_decay = 1e-4 23 | 24 | config = DefaultConfigs() 25 | -------------------------------------------------------------------------------- /dataset/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /dataset/__pycache__/aug.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/liaoxy169/pytorch-image-classification/3d38975320a3cbe21b543d72d9880d2c8f04b033/dataset/__pycache__/aug.cpython-36.pyc -------------------------------------------------------------------------------- /dataset/__pycache__/dataloader.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/liaoxy169/pytorch-image-classification/3d38975320a3cbe21b543d72d9880d2c8f04b033/dataset/__pycache__/dataloader.cpython-36.pyc -------------------------------------------------------------------------------- /dataset/aug.py: -------------------------------------------------------------------------------- 1 | from __future__ import division 2 | import cv2 3 | import numpy as np 4 | from numpy import random 5 | import math 6 | from sklearn.utils import shuffle 7 | 8 | __all__ = ['Compose','RandomHflip', 'RandomUpperCrop', 'Resize', 'UpperCrop', 'RandomBottomCrop',"RandomErasing", 9 | 'BottomCrop', 'Normalize', 'RandomSwapChannels', 'RandomRotate', 'RandomHShift',"CenterCrop","RandomVflip", 10 | 'ExpandBorder', 'RandomResizedCrop','RandomDownCrop', 'DownCrop', 'ResizedCrop',"FixRandomRotate"] 11 | 12 | def rotate_nobound(image, angle, center=None, scale=1.): 13 | (h, w) = image.shape[:2] 14 | 15 | 16 | # if the center is None, initialize it as the center of 17 | # the image 18 | if center is None: 19 | center = (w // 2, h // 2) 20 | 21 | # perform the rotation 22 | M = cv2.getRotationMatrix2D(center, angle, scale) 23 | rotated = cv2.warpAffine(image, M, (w, h)) 24 | 25 | return rotated 26 | 27 | def scale_down(src_size, size): 28 | w, h = size 29 | sw, sh = src_size 30 | if sh < h: 31 | w, h = float(w * sh) / h, sh 32 | if sw < w: 33 | w, h = sw, float(h * sw) / w 34 | return int(w), int(h) 35 | 36 | 37 | def fixed_crop(src, x0, y0, w, h, size=None): 38 | out = src[y0:y0 + h, x0:x0 + w] 39 | if size is not None and (w, h) != size: 40 | out = cv2.resize(out, (size[0], size[1]), interpolation=cv2.INTER_CUBIC) 41 | return out 42 | class FixRandomRotate(object): 43 | def __init__(self, angles=[0,90,180,270], bound=False): 44 | self.angles = angles 45 | self.bound = bound 46 | 47 | def __call__(self,img): 48 | do_rotate = random.randint(0, 4) 49 | angle=self.angles[do_rotate] 50 | if self.bound: 51 | img = rotate_bound(img, angle) 52 | else: 53 | img = rotate_nobound(img, angle) 54 | return img 55 | 56 | def center_crop(src, size): 57 | h, w = src.shape[0:2] 58 | new_w, new_h = scale_down((w, h), size) 59 | 60 | x0 = int((w - new_w) / 2) 61 | y0 = int((h - new_h) / 2) 62 | 63 | out = fixed_crop(src, x0, y0, new_w, new_h, size) 64 | return out 65 | 66 | 67 | def bottom_crop(src, size): 68 | h, w = src.shape[0:2] 69 | new_w, new_h = scale_down((w, h), size) 70 | 71 | x0 = int((w - new_w) / 2) 72 | y0 = int((h - new_h) * 0.75) 73 | 74 | out = fixed_crop(src, x0, y0, new_w, new_h, size) 75 | return out 76 | 77 | def rotate_bound(image, angle): 78 | # grab the dimensions of the image and then determine the 79 | # center 80 | h, w = image.shape[:2] 81 | 82 | (cX, cY) = (w // 2, h // 2) 83 | 84 | M = cv2.getRotationMatrix2D((cX, cY), angle, 1.0) 85 | cos = np.abs(M[0, 0]) 86 | sin = np.abs(M[0, 1]) 87 | 88 | # compute the new bounding dimensions of the image 89 | nW = int((h * sin) + (w * cos)) 90 | nH = int((h * cos) + (w * sin)) 91 | 92 | # adjust the rotation matrix to take into account translation 93 | M[0, 2] += (nW / 2) - cX 94 | M[1, 2] += (nH / 2) - cY 95 | 96 | rotated = cv2.warpAffine(image, M, (nW, nH)) 97 | 98 | return rotated 99 | 100 | 101 | class Compose(object): 102 | def __init__(self, transforms): 103 | self.transforms = transforms 104 | def __call__(self, img): 105 | for t in self.transforms: 106 | img = t(img) 107 | return img 108 | class RandomRotate(object): 109 | def __init__(self, angles, bound=False): 110 | self.angles = angles 111 | self.bound = bound 112 | 113 | def __call__(self,img): 114 | do_rotate = random.randint(0, 2) 115 | if do_rotate: 116 | angle = np.random.uniform(self.angles[0], self.angles[1]) 117 | if self.bound: 118 | img = rotate_bound(img, angle) 119 | else: 120 | img = rotate_nobound(img, angle) 121 | return img 122 | class RandomBrightness(object): 123 | def __init__(self, delta=10): 124 | assert delta >= 0 125 | assert delta <= 255 126 | self.delta = delta 127 | 128 | def __call__(self, image): 129 | if random.randint(2): 130 | delta = random.uniform(-self.delta, self.delta) 131 | image = (image + delta).clip(0.0, 255.0) 132 | # print('RandomBrightness,delta ',delta) 133 | return image 134 | 135 | 136 | class RandomContrast(object): 137 | def __init__(self, lower=0.9, upper=1.05): 138 | self.lower = lower 139 | self.upper = upper 140 | assert self.upper >= self.lower, "contrast upper must be >= lower." 141 | assert self.lower >= 0, "contrast lower must be non-negative." 142 | 143 | # expects float image 144 | def __call__(self, image): 145 | if random.randint(2): 146 | alpha = random.uniform(self.lower, self.upper) 147 | # print('contrast:', alpha) 148 | image = (image * alpha).clip(0.0,255.0) 149 | return image 150 | 151 | 152 | class RandomSaturation(object): 153 | def __init__(self, lower=0.8, upper=1.2): 154 | self.lower = lower 155 | self.upper = upper 156 | assert self.upper >= self.lower, "contrast upper must be >= lower." 157 | assert self.lower >= 0, "contrast lower must be non-negative." 158 | 159 | def __call__(self, image): 160 | if random.randint(2): 161 | alpha = random.uniform(self.lower, self.upper) 162 | image[:, :, 1] *= alpha 163 | # print('RandomSaturation,alpha',alpha) 164 | return image 165 | 166 | 167 | class RandomHue(object): 168 | def __init__(self, delta=18.0): 169 | assert delta >= 0.0 and delta <= 360.0 170 | self.delta = delta 171 | 172 | def __call__(self, image): 173 | if random.randint(2): 174 | alpha = random.uniform(-self.delta, self.delta) 175 | image[:, :, 0] += alpha 176 | image[:, :, 0][image[:, :, 0] > 360.0] -= 360.0 177 | image[:, :, 0][image[:, :, 0] < 0.0] += 360.0 178 | # print('RandomHue,alpha:', alpha) 179 | return image 180 | 181 | 182 | class ConvertColor(object): 183 | def __init__(self, current='BGR', transform='HSV'): 184 | self.transform = transform 185 | self.current = current 186 | 187 | def __call__(self, image): 188 | if self.current == 'BGR' and self.transform == 'HSV': 189 | image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) 190 | elif self.current == 'HSV' and self.transform == 'BGR': 191 | image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR) 192 | else: 193 | raise NotImplementedError 194 | return image 195 | 196 | class RandomSwapChannels(object): 197 | def __call__(self, img): 198 | if np.random.randint(2): 199 | order = np.random.permutation(3) 200 | return img[:,:,order] 201 | return img 202 | 203 | class RandomCrop(object): 204 | def __init__(self, size): 205 | self.size = size 206 | def __call__(self, image): 207 | h, w, _ = image.shape 208 | new_w, new_h = scale_down((w, h), self.size) 209 | 210 | if w == new_w: 211 | x0 = 0 212 | else: 213 | x0 = random.randint(0, w - new_w) 214 | 215 | if h == new_h: 216 | y0 = 0 217 | else: 218 | y0 = random.randint(0, h - new_h) 219 | 220 | out = fixed_crop(image, x0, y0, new_w, new_h, self.size) 221 | return out 222 | 223 | 224 | 225 | class RandomResizedCrop(object): 226 | def __init__(self, size,scale=(0.49, 1.0), ratio=(1., 1.)): 227 | self.size = size 228 | self.scale = scale 229 | self.ratio = ratio 230 | 231 | def __call__(self,img): 232 | if random.random() < 0.2: 233 | return cv2.resize(img,self.size) 234 | h, w, _ = img.shape 235 | area = h * w 236 | d=1 237 | for attempt in range(10): 238 | target_area = random.uniform(self.scale[0], self.scale[1]) * area 239 | aspect_ratio = random.uniform(self.ratio[0], self.ratio[1]) 240 | 241 | 242 | new_w = int(round(math.sqrt(target_area * aspect_ratio))) 243 | new_h = int(round(math.sqrt(target_area / aspect_ratio))) 244 | 245 | if random.random() < 0.5: 246 | new_h, new_w = new_w, new_h 247 | 248 | if new_w < w and new_h < h: 249 | x0 = random.randint(0, w - new_w) 250 | y0 = (random.randint(0, h - new_h))//d 251 | out = fixed_crop(img, x0, y0, new_w, new_h, self.size) 252 | 253 | return out 254 | 255 | # Fallback 256 | return center_crop(img, self.size) 257 | 258 | 259 | class DownCrop(): 260 | def __init__(self, size, select, scale=(0.36,0.81)): 261 | self.size = size 262 | self.scale = scale 263 | self.select = select 264 | 265 | def __call__(self,img, attr_idx): 266 | if attr_idx not in self.select: 267 | return img, attr_idx 268 | if attr_idx == 0: 269 | self.scale=(0.64,1.0) 270 | h, w, _ = img.shape 271 | area = h * w 272 | 273 | s = (self.scale[0]+self.scale[1])/2.0 274 | 275 | target_area = s * area 276 | 277 | new_w = int(round(math.sqrt(target_area))) 278 | new_h = int(round(math.sqrt(target_area))) 279 | 280 | if new_w < w and new_h < h: 281 | dw = w-new_w 282 | x0 = int(0.5*dw) 283 | y0 = h-new_h 284 | out = fixed_crop(img, x0, y0, new_w, new_h, self.size) 285 | return out, attr_idx 286 | 287 | # Fallback 288 | return center_crop(img, self.size), attr_idx 289 | 290 | 291 | class ResizedCrop(object): 292 | def __init__(self, size, select,scale=(0.64, 1.0), ratio=(3. / 4., 4. / 3.)): 293 | self.size = size 294 | self.scale = scale 295 | self.ratio = ratio 296 | self.select = select 297 | 298 | def __call__(self,img, attr_idx): 299 | if attr_idx not in self.select: 300 | return img, attr_idx 301 | h, w, _ = img.shape 302 | area = h * w 303 | d=1 304 | if attr_idx == 2: 305 | self.scale=(0.36,0.81) 306 | d=2 307 | if attr_idx == 0: 308 | self.scale=(0.81,1.0) 309 | 310 | target_area = (self.scale[0]+self.scale[1])/2.0 * area 311 | # aspect_ratio = random.uniform(self.ratio[0], self.ratio[1]) 312 | 313 | 314 | new_w = int(round(math.sqrt(target_area))) 315 | new_h = int(round(math.sqrt(target_area))) 316 | 317 | # if random.random() < 0.5: 318 | # new_h, new_w = new_w, new_h 319 | 320 | if new_w < w and new_h < h: 321 | x0 = (w - new_w)//2 322 | y0 = (h - new_h)//d//2 323 | out = fixed_crop(img, x0, y0, new_w, new_h, self.size) 324 | # cv2.imshow('{}_img'.format(idx2attr_map[attr_idx]), img) 325 | # cv2.imshow('{}_crop'.format(idx2attr_map[attr_idx]), out) 326 | # 327 | # cv2.waitKey(0) 328 | return out, attr_idx 329 | 330 | # Fallback 331 | return center_crop(img, self.size), attr_idx 332 | 333 | class RandomHflip(object): 334 | def __call__(self, image): 335 | if random.randint(2): 336 | return cv2.flip(image, 1) 337 | else: 338 | return image 339 | class RandomVflip(object): 340 | def __call__(self, image): 341 | if random.randint(2): 342 | return cv2.flip(image, 0) 343 | else: 344 | return image 345 | 346 | 347 | class Hflip(object): 348 | def __init__(self,doHflip): 349 | self.doHflip = doHflip 350 | 351 | def __call__(self, image): 352 | if self.doHflip: 353 | return cv2.flip(image, 1) 354 | else: 355 | return image 356 | 357 | 358 | class CenterCrop(object): 359 | def __init__(self, size): 360 | self.size = size 361 | 362 | def __call__(self, image): 363 | return center_crop(image, self.size) 364 | 365 | class UpperCrop(): 366 | def __init__(self, size, scale=(0.09, 0.64)): 367 | self.size = size 368 | self.scale = scale 369 | 370 | def __call__(self,img): 371 | h, w, _ = img.shape 372 | area = h * w 373 | 374 | s = (self.scale[0]+self.scale[1])/2.0 375 | 376 | target_area = s * area 377 | 378 | new_w = int(round(math.sqrt(target_area))) 379 | new_h = int(round(math.sqrt(target_area))) 380 | 381 | if new_w < w and new_h < h: 382 | dw = w-new_w 383 | x0 = int(0.5*dw) 384 | y0 = 0 385 | out = fixed_crop(img, x0, y0, new_w, new_h, self.size) 386 | return out 387 | 388 | # Fallback 389 | return center_crop(img, self.size) 390 | 391 | 392 | 393 | class RandomUpperCrop(object): 394 | def __init__(self, size, select, scale=(0.09, 0.64), ratio=(3. / 4., 4. / 3.)): 395 | self.size = size 396 | self.scale = scale 397 | self.ratio = ratio 398 | self.select = select 399 | 400 | def __call__(self,img, attr_idx): 401 | if random.random() < 0.2: 402 | return img, attr_idx 403 | if attr_idx not in self.select: 404 | return img, attr_idx 405 | 406 | h, w, _ = img.shape 407 | area = h * w 408 | for attempt in range(10): 409 | s = random.uniform(self.scale[0], self.scale[1]) 410 | d = 0.1 + (0.3 - 0.1) / (self.scale[1] - self.scale[0]) * (s - self.scale[0]) 411 | target_area = s * area 412 | aspect_ratio = random.uniform(self.ratio[0], self.ratio[1]) 413 | new_w = int(round(math.sqrt(target_area * aspect_ratio))) 414 | new_h = int(round(math.sqrt(target_area / aspect_ratio))) 415 | 416 | 417 | # new_w = int(round(math.sqrt(target_area))) 418 | # new_h = int(round(math.sqrt(target_area))) 419 | 420 | if new_w < w and new_h < h: 421 | dw = w-new_w 422 | x0 = random.randint(int((0.5-d)*dw), int((0.5+d)*dw)+1) 423 | y0 = (random.randint(0, h - new_h))//10 424 | out = fixed_crop(img, x0, y0, new_w, new_h, self.size) 425 | return out, attr_idx 426 | 427 | # Fallback 428 | return center_crop(img, self.size), attr_idx 429 | class RandomDownCrop(object): 430 | def __init__(self, size, select, scale=(0.36, 0.81), ratio=(3. / 4., 4. / 3.)): 431 | self.size = size 432 | self.scale = scale 433 | self.ratio = ratio 434 | self.select = select 435 | 436 | def __call__(self,img, attr_idx): 437 | if random.random() < 0.2: 438 | return img, attr_idx 439 | if attr_idx not in self.select: 440 | return img, attr_idx 441 | if attr_idx == 0: 442 | self.scale=(0.64,1.0) 443 | 444 | h, w, _ = img.shape 445 | area = h * w 446 | for attempt in range(10): 447 | s = random.uniform(self.scale[0], self.scale[1]) 448 | d = 0.1 + (0.3 - 0.1) / (self.scale[1] - self.scale[0]) * (s - self.scale[0]) 449 | target_area = s * area 450 | aspect_ratio = random.uniform(self.ratio[0], self.ratio[1]) 451 | new_w = int(round(math.sqrt(target_area * aspect_ratio))) 452 | new_h = int(round(math.sqrt(target_area / aspect_ratio))) 453 | # 454 | # new_w = int(round(math.sqrt(target_area))) 455 | # new_h = int(round(math.sqrt(target_area))) 456 | 457 | if new_w < w and new_h < h: 458 | dw = w-new_w 459 | x0 = random.randint(int((0.5-d)*dw), int((0.5+d)*dw)+1) 460 | y0 = (random.randint((h - new_h)*9//10, h - new_h)) 461 | out = fixed_crop(img, x0, y0, new_w, new_h, self.size) 462 | 463 | # cv2.imshow('{}_img'.format(idx2attr_map[attr_idx]), img) 464 | # cv2.imshow('{}_crop'.format(idx2attr_map[attr_idx]), out) 465 | # 466 | # cv2.waitKey(0) 467 | 468 | return out, attr_idx 469 | 470 | # Fallback 471 | return center_crop(img, self.size), attr_idx 472 | 473 | class RandomHShift(object): 474 | def __init__(self, select, scale=(0.0, 0.2)): 475 | self.scale = scale 476 | self.select = select 477 | 478 | def __call__(self,img, attr_idx): 479 | if attr_idx not in self.select: 480 | return img, attr_idx 481 | do_shift_crop = random.randint(0, 2) 482 | if do_shift_crop: 483 | h, w, _ = img.shape 484 | min_shift = int(w*self.scale[0]) 485 | max_shift = int(w*self.scale[1]) 486 | shift_idx = random.randint(min_shift, max_shift) 487 | direction = random.randint(0,2) 488 | if direction: 489 | right_part = img[:, -shift_idx:, :] 490 | left_part = img[:, :-shift_idx, :] 491 | else: 492 | left_part = img[:, :shift_idx, :] 493 | right_part = img[:, shift_idx:, :] 494 | img = np.concatenate((right_part, left_part), axis=1) 495 | 496 | # Fallback 497 | return img, attr_idx 498 | 499 | 500 | class RandomBottomCrop(object): 501 | def __init__(self, size, select, scale=(0.4, 0.8)): 502 | self.size = size 503 | self.scale = scale 504 | self.select = select 505 | 506 | def __call__(self,img, attr_idx): 507 | if attr_idx not in self.select: 508 | return img, attr_idx 509 | 510 | h, w, _ = img.shape 511 | area = h * w 512 | for attempt in range(10): 513 | s = random.uniform(self.scale[0], self.scale[1]) 514 | d = 0.25 + (0.45 - 0.25) / (self.scale[1] - self.scale[0]) * (s - self.scale[0]) 515 | target_area = s * area 516 | 517 | new_w = int(round(math.sqrt(target_area))) 518 | new_h = int(round(math.sqrt(target_area))) 519 | 520 | if new_w < w and new_h < h: 521 | dw = w-new_w 522 | dh = h - new_h 523 | x0 = random.randint(int((0.5-d)*dw), min(int((0.5+d)*dw)+1,dw)) 524 | y0 = (random.randint(max(0,int(0.8*dh)-1), dh)) 525 | out = fixed_crop(img, x0, y0, new_w, new_h, self.size) 526 | return out, attr_idx 527 | 528 | # Fallback 529 | return bottom_crop(img, self.size), attr_idx 530 | 531 | 532 | class BottomCrop(): 533 | def __init__(self, size, select, scale=(0.4, 0.8)): 534 | self.size = size 535 | self.scale = scale 536 | self.select = select 537 | 538 | def __call__(self,img, attr_idx): 539 | if attr_idx not in self.select: 540 | return img, attr_idx 541 | 542 | h, w, _ = img.shape 543 | area = h * w 544 | 545 | s = (self.scale[0]+self.scale[1])/3.*2. 546 | 547 | target_area = s * area 548 | 549 | new_w = int(round(math.sqrt(target_area))) 550 | new_h = int(round(math.sqrt(target_area))) 551 | 552 | if new_w < w and new_h < h: 553 | dw = w-new_w 554 | dh = h-new_h 555 | x0 = int(0.5*dw) 556 | y0 = int(0.9*dh) 557 | out = fixed_crop(img, x0, y0, new_w, new_h, self.size) 558 | return out, attr_idx 559 | 560 | # Fallback 561 | return bottom_crop(img, self.size), attr_idx 562 | 563 | 564 | 565 | class Resize(object): 566 | def __init__(self, size, inter=cv2.INTER_CUBIC): 567 | self.size = size 568 | self.inter = inter 569 | 570 | def __call__(self, image): 571 | return cv2.resize(image, (self.size[0], self.size[0]), interpolation=self.inter) 572 | 573 | class ExpandBorder(object): 574 | def __init__(self, mode='constant', value=255, size=(336,336), resize=False): 575 | self.mode = mode 576 | self.value = value 577 | self.resize = resize 578 | self.size = size 579 | 580 | def __call__(self, image): 581 | h, w, _ = image.shape 582 | if h > w: 583 | pad1 = (h-w)//2 584 | pad2 = h - w - pad1 585 | if self.mode == 'constant': 586 | image = np.pad(image, ((0, 0), (pad1, pad2), (0, 0)), 587 | self.mode, constant_values=self.value) 588 | else: 589 | image = np.pad(image,((0,0), (pad1, pad2),(0,0)), self.mode) 590 | elif h < w: 591 | pad1 = (w-h)//2 592 | pad2 = w-h - pad1 593 | if self.mode == 'constant': 594 | image = np.pad(image, ((pad1, pad2),(0, 0), (0, 0)), 595 | self.mode,constant_values=self.value) 596 | else: 597 | image = np.pad(image, ((pad1, pad2), (0, 0), (0, 0)),self.mode) 598 | if self.resize: 599 | image = cv2.resize(image, (self.size[0], self.size[0]),interpolation=cv2.INTER_LINEAR) 600 | return image 601 | class AstypeToInt(): 602 | def __call__(self, image, attr_idx): 603 | return image.clip(0,255.0).astype(np.uint8), attr_idx 604 | 605 | class AstypeToFloat(): 606 | def __call__(self, image, attr_idx): 607 | return image.astype(np.float32), attr_idx 608 | 609 | import matplotlib.pyplot as plt 610 | class Normalize(object): 611 | def __init__(self,mean, std): 612 | ''' 613 | :param mean: RGB order 614 | :param std: RGB order 615 | ''' 616 | self.mean = np.array(mean).reshape(3,1,1) 617 | self.std = np.array(std).reshape(3,1,1) 618 | def __call__(self, image): 619 | ''' 620 | :param image: (H,W,3) RGB 621 | :return: 622 | ''' 623 | # plt.figure(1) 624 | # plt.imshow(image) 625 | # plt.show() 626 | return (image.transpose((2, 0, 1)) / 255. - self.mean) / self.std 627 | 628 | class RandomErasing(object): 629 | def __init__(self, select,EPSILON=0.5,sl=0.02, sh=0.09, r1=0.3, mean=[0.485, 0.456, 0.406]): 630 | self.EPSILON = EPSILON 631 | self.mean = mean 632 | self.sl = sl 633 | self.sh = sh 634 | self.r1 = r1 635 | self.select = select 636 | 637 | def __call__(self, img,attr_idx): 638 | if attr_idx not in self.select: 639 | return img,attr_idx 640 | 641 | if random.uniform(0, 1) > self.EPSILON: 642 | return img,attr_idx 643 | 644 | for attempt in range(100): 645 | area = img.shape[1] * img.shape[2] 646 | 647 | target_area = random.uniform(self.sl, self.sh) * area 648 | aspect_ratio = random.uniform(self.r1, 1 / self.r1) 649 | 650 | h = int(round(math.sqrt(target_area * aspect_ratio))) 651 | w = int(round(math.sqrt(target_area / aspect_ratio))) 652 | 653 | if w <= img.shape[2] and h <= img.shape[1]: 654 | x1 = random.randint(0, img.shape[1] - h) 655 | y1 = random.randint(0, img.shape[2] - w) 656 | if img.shape[0] == 3: 657 | # img[0, x1:x1+h, y1:y1+w] = random.uniform(0, 1) 658 | # img[1, x1:x1+h, y1:y1+w] = random.uniform(0, 1) 659 | # img[2, x1:x1+h, y1:y1+w] = random.uniform(0, 1) 660 | img[0, x1:x1 + h, y1:y1 + w] = self.mean[0] 661 | img[1, x1:x1 + h, y1:y1 + w] = self.mean[1] 662 | img[2, x1:x1 + h, y1:y1 + w] = self.mean[2] 663 | # img[:, x1:x1+h, y1:y1+w] = torch.from_numpy(np.random.rand(3, h, w)) 664 | else: 665 | img[0, x1:x1 + h, y1:y1 + w] = self.mean[1] 666 | # img[0, x1:x1+h, y1:y1+w] = torch.from_numpy(np.random.rand(1, h, w)) 667 | return img,attr_idx 668 | 669 | return img,attr_idx 670 | 671 | if __name__ == '__main__': 672 | import matplotlib.pyplot as plt 673 | 674 | 675 | class FSAug(object): 676 | def __init__(self): 677 | self.augment = Compose([ 678 | AstypeToFloat(), 679 | # RandomHShift(scale=(0.,0.2),select=range(8)), 680 | # RandomRotate(angles=(-20., 20.), bound=True), 681 | ExpandBorder(select=range(8), mode='symmetric'),# symmetric 682 | # Resize(size=(336, 336), select=[ 2, 7]), 683 | AstypeToInt() 684 | ]) 685 | 686 | def __call__(self, spct,attr_idx): 687 | return self.augment(spct,attr_idx) 688 | 689 | 690 | trans = FSAug() 691 | 692 | img_path = '/media/gserver/data/FashionAI/round2/train/Images/coat_length_labels/0b6b4a2146fc8616a19fcf2026d61d50.jpg' 693 | img = cv2.cvtColor(cv2.imread(img_path),cv2.COLOR_BGR2RGB) 694 | img_trans,_ = trans(img,5) 695 | # img_trans2,_ = trans(img,6) 696 | 697 | plt.figure() 698 | plt.subplot(221) 699 | plt.imshow(img) 700 | 701 | plt.subplot(222) 702 | plt.imshow(img_trans) 703 | 704 | # plt.subplot(223) 705 | # plt.imshow(img_trans2) 706 | # plt.imshow(img_trans2) 707 | plt.show() 708 | -------------------------------------------------------------------------------- /dataset/dataloader.py: -------------------------------------------------------------------------------- 1 | from torch.utils.data import Dataset 2 | from torchvision import transforms as T 3 | from config import config 4 | from PIL import Image 5 | from dataset.aug import * 6 | from itertools import chain 7 | from glob import glob 8 | from tqdm import tqdm 9 | import random 10 | import numpy as np 11 | import pandas as pd 12 | import os 13 | import cv2 14 | import torch 15 | 16 | #1.set random seed 17 | random.seed(config.seed) 18 | np.random.seed(config.seed) 19 | torch.manual_seed(config.seed) 20 | torch.cuda.manual_seed_all(config.seed) 21 | 22 | #2.define dataset 23 | class ChaojieDataset(Dataset): 24 | def __init__(self,label_list,transforms=None,train=True,test=False): 25 | self.test = test 26 | self.train = train 27 | imgs = [] 28 | if self.test: 29 | for index,row in label_list.iterrows(): 30 | imgs.append((row["filename"])) 31 | self.imgs = imgs 32 | else: 33 | for index,row in label_list.iterrows(): 34 | imgs.append((row["filename"],row["label"])) 35 | self.imgs = imgs 36 | if transforms is None: 37 | if self.test or not self.train: 38 | self.transforms = T.Compose([ 39 | T.Resize((config.img_weight,config.img_height)), 40 | T.ToTensor(), 41 | T.Normalize(mean = [0.485,0.456,0.406], 42 | std = [0.229,0.224,0.225])]) 43 | else: 44 | self.transforms = T.Compose([ 45 | T.Resize((config.img_weight,config.img_height)), 46 | T.RandomRotation(30), 47 | T.RandomHorizontalFlip(), 48 | T.RandomVerticalFlip(), 49 | T.RandomAffine(45), 50 | T.ToTensor(), 51 | T.Normalize(mean = [0.485,0.456,0.406], 52 | std = [0.229,0.224,0.225])]) 53 | else: 54 | self.transforms = transforms 55 | def __getitem__(self,index): 56 | if self.test: 57 | filename = self.imgs[index] 58 | #img = cv2.imread(filename) 59 | #img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB) 60 | img = Image.open(filename) 61 | img = self.transforms(img) 62 | return img,filename 63 | else: 64 | filename,label = self.imgs[index] 65 | #img = cv2.imread(filename) 66 | #img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB) 67 | img = Image.open(filename) 68 | img = self.transforms(img) 69 | return img,label 70 | def __len__(self): 71 | return len(self.imgs) 72 | 73 | def collate_fn(batch): 74 | imgs = [] 75 | label = [] 76 | for sample in batch: 77 | imgs.append(sample[0]) 78 | label.append(sample[1]) 79 | 80 | return torch.stack(imgs, 0), \ 81 | label 82 | 83 | def get_files(root,mode): 84 | #for test 85 | if mode == "test": 86 | files = [] 87 | for img in os.listdir(root): 88 | files.append(root + img) 89 | files = pd.DataFrame({"filename":files}) 90 | return files 91 | elif mode != "test": 92 | #for train and val 93 | all_data_path,labels = [],[] 94 | image_folders = list(map(lambda x:root+x,os.listdir(root))) 95 | all_images = list(chain.from_iterable(list(map(lambda x:glob(x+"/*"),image_folders)))) 96 | print("loading train dataset") 97 | for file in tqdm(all_images): 98 | all_data_path.append(file) 99 | labels.append(int(file.split("/")[-2])) 100 | all_files = pd.DataFrame({"filename":all_data_path,"label":labels}) 101 | return all_files 102 | else: 103 | print("check the mode please!") 104 | 105 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import time 4 | import json 5 | import torch 6 | import torchvision 7 | import numpy as np 8 | import pandas as pd 9 | import warnings 10 | from datetime import datetime 11 | from torch import nn,optim 12 | from config import config 13 | from collections import OrderedDict 14 | from torch.autograd import Variable 15 | from torch.utils.data import DataLoader 16 | from dataset.dataloader import * 17 | from sklearn.model_selection import train_test_split,StratifiedKFold 18 | from timeit import default_timer as timer 19 | from models.model import * 20 | from utils import * 21 | from IPython import embed 22 | #1. set random.seed and cudnn performance 23 | random.seed(config.seed) 24 | np.random.seed(config.seed) 25 | torch.manual_seed(config.seed) 26 | torch.cuda.manual_seed_all(config.seed) 27 | os.environ["CUDA_VISIBLE_DEVICES"] = config.gpus 28 | torch.backends.cudnn.benchmark = True 29 | warnings.filterwarnings('ignore') 30 | 31 | #2. evaluate func 32 | def evaluate(val_loader,model,criterion,epoch): 33 | #2.1 define meters 34 | losses = AverageMeter() 35 | top1 = AverageMeter() 36 | #progress bar 37 | val_progressor = ProgressBar(mode="Val ",epoch=epoch,total_epoch=config.epochs,model_name=config.model_name,total=len(val_loader)) 38 | #2.2 switch to evaluate mode and confirm model has been transfered to cuda 39 | model.cuda() 40 | #model.eval() 41 | with torch.no_grad(): 42 | for i,(input,target) in enumerate(val_loader): 43 | val_progressor.current = i 44 | input = Variable(input).cuda() 45 | target = Variable(torch.from_numpy(np.array(target)).long()).cuda() 46 | #target = Variable(target).cuda() 47 | #2.2.1 compute output 48 | output = model(input) 49 | loss = criterion(output,target) 50 | 51 | #2.2.2 measure accuracy and record loss 52 | precision1,precision2 = accuracy(output,target,topk=(1,2)) 53 | losses.update(loss.item(),input.size(0)) 54 | top1.update(precision1[0],input.size(0)) 55 | val_progressor.current_loss = losses.avg 56 | val_progressor.current_top1 = top1.avg 57 | val_progressor() 58 | val_progressor.done() 59 | return [losses.avg,top1.avg] 60 | 61 | #3. test model on public dataset and save the probability matrix 62 | def test(test_loader,model,folds): 63 | #3.1 confirm the model converted to cuda 64 | csv_map = OrderedDict({"filename":[],"probability":[]}) 65 | model.cuda() 66 | model.eval() 67 | for i,(input,filepath) in enumerate(tqdm(test_loader)): 68 | #3.2 change everything to cuda and get only basename 69 | filepath = [os.path.basename(x) for x in filepath] 70 | with torch.no_grad(): 71 | image_var = Variable(input).cuda() 72 | #3.3.output 73 | #print(filepath) 74 | #print(input,input.shape) 75 | y_pred = model(image_var) 76 | print(y_pred.shape) 77 | smax = nn.Softmax(1) 78 | smax_out = smax(y_pred) 79 | #3.4 save probability to csv files 80 | csv_map["filename"].extend(filepath) 81 | for output in smax_out: 82 | prob = ";".join([str(i) for i in output.data.tolist()]) 83 | csv_map["probability"].append(prob) 84 | result = pd.DataFrame(csv_map) 85 | result["probability"] = result["probability"].map(lambda x : [float(i) for i in x.split(";")]) 86 | result.to_csv("./submit/{}_submission.csv" .format(config.model_name + "_" + str(folds)),index=False,header = None) 87 | 88 | #4. more details to build main function 89 | def main(): 90 | fold = 0 91 | #4.1 mkdirs 92 | if not os.path.exists(config.submit): 93 | os.mkdir(config.submit) 94 | if not os.path.exists(config.weights): 95 | os.mkdir(config.weights) 96 | if not os.path.exists(config.best_models): 97 | os.mkdir(config.best_models) 98 | if not os.path.exists(config.logs): 99 | os.mkdir(config.logs) 100 | if not os.path.exists(config.weights + config.model_name + os.sep +str(fold) + os.sep): 101 | os.makedirs(config.weights + config.model_name + os.sep +str(fold) + os.sep) 102 | if not os.path.exists(config.best_models + config.model_name + os.sep +str(fold) + os.sep): 103 | os.makedirs(config.best_models + config.model_name + os.sep +str(fold) + os.sep) 104 | #4.2 get model and optimizer 105 | model = get_net() 106 | #model = torch.nn.DataParallel(model) 107 | model.cuda() 108 | 109 | #optimizer = optim.SGD(model.parameters(),lr = config.lr,momentum=0.9,weight_decay=config.weight_decay) 110 | optimizer = optim.Adam(model.parameters(),lr = config.lr,amsgrad=True,weight_decay=config.weight_decay) 111 | criterion = nn.CrossEntropyLoss().cuda() 112 | #4.3 some parameters for K-fold and restart model 113 | start_epoch = 0 114 | best_precision1 = 0 115 | best_precision_save = 0 116 | resume = False 117 | 118 | #4.4 restart the training process 119 | if resume: 120 | checkpoint = torch.load(config.best_models + str(fold) + "/model_best.pth.tar") 121 | start_epoch = checkpoint["epoch"] 122 | fold = checkpoint["fold"] 123 | best_precision1 = checkpoint["best_precision1"] 124 | model.load_state_dict(checkpoint["state_dict"]) 125 | optimizer.load_state_dict(checkpoint["optimizer"]) 126 | 127 | #4.5 get files and split for K-fold dataset 128 | #4.5.1 read files 129 | train_data_list = get_files(config.train_data,"train") 130 | val_data_list = get_files(config.val_data,"val") 131 | #test_files = get_files(config.test_data,"test") 132 | 133 | """ 134 | #4.5.2 split 135 | split_fold = StratifiedKFold(n_splits=3) 136 | folds_indexes = split_fold.split(X=origin_files["filename"],y=origin_files["label"]) 137 | folds_indexes = np.array(list(folds_indexes)) 138 | fold_index = folds_indexes[fold] 139 | 140 | #4.5.3 using fold index to split for train data and val data 141 | train_data_list = pd.concat([origin_files["filename"][fold_index[0]],origin_files["label"][fold_index[0]]],axis=1) 142 | val_data_list = pd.concat([origin_files["filename"][fold_index[1]],origin_files["label"][fold_index[1]]],axis=1) 143 | """ 144 | #train_data_list,val_data_list = train_test_split(origin_files,test_size = 0.1,stratify=origin_files["label"]) 145 | #4.5.4 load dataset 146 | train_dataloader = DataLoader(ChaojieDataset(train_data_list),batch_size=config.batch_size,shuffle=True,collate_fn=collate_fn,pin_memory=True) 147 | val_dataloader = DataLoader(ChaojieDataset(val_data_list,train=False),batch_size=config.batch_size * 2,shuffle=True,collate_fn=collate_fn,pin_memory=False) 148 | #test_dataloader = DataLoader(ChaojieDataset(test_files,test=True),batch_size=1,shuffle=False,pin_memory=False) 149 | #scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,"max",verbose=1,patience=3) 150 | scheduler = optim.lr_scheduler.StepLR(optimizer,step_size = 10,gamma=0.1) 151 | #4.5.5.1 define metrics 152 | train_losses = AverageMeter() 153 | train_top1 = AverageMeter() 154 | valid_loss = [np.inf,0,0] 155 | model.train() 156 | 157 | #4.5.5 train 158 | start = timer() 159 | for epoch in range(start_epoch,config.epochs): 160 | scheduler.step(epoch) 161 | train_progressor = ProgressBar(mode="Train",epoch=epoch,total_epoch=config.epochs,model_name=config.model_name,total=len(train_dataloader)) 162 | # train 163 | #global iter 164 | for iter,(input,target) in enumerate(train_dataloader): 165 | #4.5.5 switch to continue train process 166 | train_progressor.current = iter 167 | model.train() 168 | input = Variable(input).cuda() 169 | target = Variable(torch.from_numpy(np.array(target)).long()).cuda() 170 | #target = Variable(target).cuda() 171 | output = model(input) 172 | loss = criterion(output,target) 173 | 174 | precision1_train,precision2_train = accuracy(output,target,topk=(1,2)) 175 | train_losses.update(loss.item(),input.size(0)) 176 | train_top1.update(precision1_train[0],input.size(0)) 177 | train_progressor.current_loss = train_losses.avg 178 | train_progressor.current_top1 = train_top1.avg 179 | #backward 180 | optimizer.zero_grad() 181 | loss.backward() 182 | optimizer.step() 183 | train_progressor() 184 | train_progressor.done() 185 | #evaluate 186 | lr = get_learning_rate(optimizer) 187 | #evaluate every half epoch 188 | valid_loss = evaluate(val_dataloader,model,criterion,epoch) 189 | is_best = valid_loss[1] > best_precision1 190 | best_precision1 = max(valid_loss[1],best_precision1) 191 | try: 192 | best_precision_save = best_precision1.cpu().data.numpy() 193 | except: 194 | pass 195 | save_checkpoint({ 196 | "epoch":epoch + 1, 197 | "model_name":config.model_name, 198 | "state_dict":model.state_dict(), 199 | "best_precision1":best_precision1, 200 | "optimizer":optimizer.state_dict(), 201 | "fold":fold, 202 | "valid_loss":valid_loss, 203 | },is_best,fold) 204 | 205 | if __name__ =="__main__": 206 | main() 207 | 208 | 209 | 210 | 211 | 212 | 213 | 214 | 215 | 216 | 217 | 218 | 219 | 220 | 221 | 222 | 223 | 224 | 225 | 226 | 227 | 228 | -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/liaoxy169/pytorch-image-classification/3d38975320a3cbe21b543d72d9880d2c8f04b033/models/__init__.py -------------------------------------------------------------------------------- /models/model.py: -------------------------------------------------------------------------------- 1 | import torchvision 2 | import torch.nn.functional as F 3 | from torch import nn 4 | from config import config 5 | 6 | def generate_model(): 7 | class DenseModel(nn.Module): 8 | def __init__(self, pretrained_model): 9 | super(DenseModel, self).__init__() 10 | self.classifier = nn.Linear(pretrained_model.classifier.in_features, config.num_classes) 11 | 12 | for m in self.modules(): 13 | if isinstance(m, nn.Conv2d): 14 | nn.init.kaiming_normal(m.weight) 15 | elif isinstance(m, nn.BatchNorm2d): 16 | m.weight.data.fill_(1) 17 | m.bias.data.zero_() 18 | elif isinstance(m, nn.Linear): 19 | m.bias.data.zero_() 20 | 21 | self.features = pretrained_model.features 22 | self.layer1 = pretrained_model.features._modules['denseblock1'] 23 | self.layer2 = pretrained_model.features._modules['denseblock2'] 24 | self.layer3 = pretrained_model.features._modules['denseblock3'] 25 | self.layer4 = pretrained_model.features._modules['denseblock4'] 26 | 27 | def forward(self, x): 28 | features = self.features(x) 29 | out = F.relu(features, inplace=True) 30 | out = F.avg_pool2d(out, kernel_size=8).view(features.size(0), -1) 31 | out = F.sigmoid(self.classifier(out)) 32 | return out 33 | 34 | return DenseModel(torchvision.models.densenet169(pretrained=True)) 35 | 36 | def get_net(): 37 | #return MyModel(torchvision.models.resnet101(pretrained = True)) 38 | model = torchvision.models.resnet152(pretrained = True) 39 | #for param in model.parameters(): 40 | # param.requires_grad = False 41 | model.avgpool = nn.AdaptiveAvgPool2d(1) 42 | model.fc = nn.Linear(2048,config.num_classes) 43 | return model 44 | 45 | -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import time 4 | import json 5 | import torch 6 | import torchvision 7 | import numpy as np 8 | import pandas as pd 9 | import warnings 10 | from datetime import datetime 11 | from torch import nn,optim 12 | from config import config 13 | from collections import OrderedDict 14 | from torch.autograd import Variable 15 | from torch.utils.data import DataLoader 16 | from dataset.dataloader import * 17 | from sklearn.model_selection import train_test_split,StratifiedKFold 18 | from timeit import default_timer as timer 19 | from models.model import * 20 | from utils import * 21 | from IPython import embed 22 | #1. set random.seed and cudnn performance 23 | random.seed(config.seed) 24 | np.random.seed(config.seed) 25 | torch.manual_seed(config.seed) 26 | torch.cuda.manual_seed_all(config.seed) 27 | os.environ["CUDA_VISIBLE_DEVICES"] = config.gpus 28 | torch.backends.cudnn.benchmark = True 29 | warnings.filterwarnings('ignore') 30 | 31 | #3. test model on public dataset and save the probability matrix 32 | def test(test_loader,model,folds): 33 | #3.1 confirm the model converted to cuda 34 | test_labels = open("%s.csv"%config.model_name,"w") 35 | model.cuda() 36 | model.eval() 37 | tta = False 38 | for i,(input,filepath) in enumerate(tqdm(test_loader)): 39 | #3.2 change everything to cuda and get only basename 40 | filepath = [os.path.basename(x) for x in filepath] 41 | with torch.no_grad(): 42 | image_var = Variable(input).cuda() 43 | #3.3.output 44 | #print(filepath) 45 | #print(input,input.shape) 46 | if tta == False: 47 | y_pred = model(image_var) 48 | smax = nn.Softmax(1) 49 | smax_out = smax(y_pred) 50 | label = np.argmax(smax_out.cpu().data.numpy()) 51 | test_labels.write(filepath[0]+","+str(label)+"\n") 52 | 53 | #4. more details to build main function 54 | def main(): 55 | fold = 0 56 | #4.1 mkdirs 57 | if not os.path.exists(config.submit): 58 | os.mkdir(config.submit) 59 | if not os.path.exists(config.weights): 60 | os.mkdir(config.weights) 61 | if not os.path.exists(config.best_models): 62 | os.mkdir(config.best_models) 63 | if not os.path.exists(config.logs): 64 | os.mkdir(config.logs) 65 | if not os.path.exists(config.weights + config.model_name + os.sep +str(fold) + os.sep): 66 | os.makedirs(config.weights + config.model_name + os.sep +str(fold) + os.sep) 67 | if not os.path.exists(config.best_models + config.model_name + os.sep +str(fold) + os.sep): 68 | os.makedirs(config.best_models + config.model_name + os.sep +str(fold) + os.sep) 69 | #4.2 get model and optimizer 70 | model = get_net() 71 | #model = torch.nn.DataParallel(model) 72 | model.cuda() 73 | #4.5 get files and split for K-fold dataset 74 | test_files = get_files(config.test_data,"test") 75 | #4.5.4 load dataset 76 | test_dataloader = DataLoader(ChaojieDataset(test_files,test=True),batch_size=1,shuffle=False,pin_memory=False) 77 | best_model = torch.load("checkpoints/best_model/%s/0/model_best.pth.tar"%config.model_name) 78 | model.load_state_dict(best_model["state_dict"]) 79 | test(test_dataloader,model,fold) 80 | 81 | if __name__ =="__main__": 82 | main() 83 | 84 | 85 | 86 | 87 | 88 | 89 | 90 | 91 | 92 | 93 | 94 | 95 | 96 | 97 | 98 | 99 | 100 | 101 | 102 | 103 | 104 | -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .utils import * 2 | from .progress_bar import * -------------------------------------------------------------------------------- /utils/progress_bar.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import re 3 | class ProgressBar(object): 4 | DEFAULT = "Progress: %(bar)s %(percent)3d%%" 5 | def __init__(self,mode,epoch=None,total_epoch=None,current_loss=None,current_top1=None,model_name=None,total=None,current=None,width = 50,symbol = ">",output=sys.stderr): 6 | assert len(symbol) == 1 7 | 8 | self.mode = mode 9 | self.total = total 10 | self.symbol = symbol 11 | self.output = output 12 | self.width = width 13 | self.current = current 14 | self.epoch = epoch 15 | self.total_epoch = total_epoch 16 | self.current_loss = current_loss 17 | self.current_top1 = current_top1 18 | self.model_name = model_name 19 | 20 | def __call__(self): 21 | percent = self.current / float(self.total) 22 | size = int(self.width * percent) 23 | bar = "[" + self.symbol * size + " " * (self.width - size) + "]" 24 | 25 | args = { 26 | "mode":self.mode, 27 | "total": self.total, 28 | "bar" : bar, 29 | "current": self.current, 30 | "percent": percent * 100, 31 | "current_loss":self.current_loss, 32 | "current_top1":self.current_top1, 33 | "epoch":self.epoch + 1, 34 | "epochs":self.total_epoch 35 | } 36 | message = "\033[1;32;40m%(mode)s Epoch: %(epoch)d/%(epochs)d %(bar)s\033[0m [Current: Loss %(current_loss)f Top1: %(current_top1)f ] %(current)d/%(total)d \033[1;32;40m[ %(percent)3d%% ]\033[0m" %args 37 | self.write_message = "%(mode)s Epoch: %(epoch)d/%(epochs)d %(bar)s [Current: Loss %(current_loss)f Top1: %(current_top1)f ] %(current)d/%(total)d [ %(percent)3d%% ]" %args 38 | print("\r" + message,file=self.output,end="") 39 | 40 | 41 | def done(self): 42 | self.current = self.total 43 | self() 44 | print("",file=self.output) 45 | with open("./logs/%s.txt"%self.model_name,"a") as f: 46 | print(self.write_message,file=f) 47 | if __name__ == "__main__": 48 | 49 | from time import sleep 50 | progress = ProgressBar("Train",total_epoch=150,model_name="resnet159") 51 | for i in range(150): 52 | progress.total = 50 53 | progress.epoch = i 54 | progress.current_loss = 0.15 55 | progress.current_top1 = 0.45 56 | for x in range(50): 57 | progress.current = x 58 | progress() 59 | sleep(0.1) 60 | progress.done() 61 | -------------------------------------------------------------------------------- /utils/utils.py: -------------------------------------------------------------------------------- 1 | import shutil 2 | import torch 3 | import sys 4 | import os 5 | from config import config 6 | 7 | def save_checkpoint(state, is_best,fold): 8 | filename = config.weights + config.model_name + os.sep +str(fold) + os.sep + "_checkpoint.pth.tar" 9 | torch.save(state, filename) 10 | if is_best: 11 | message = config.best_models + config.model_name+ os.sep +str(fold) + os.sep + 'model_best.pth.tar' 12 | print("Get Better top1 : %s saving weights to %s"%(state["best_precision1"],message)) 13 | with open("./logs/%s.txt"%config.model_name,"a") as f: 14 | print("Get Better top1 : %s saving weights to %s"%(state["best_precision1"],message),file=f) 15 | shutil.copyfile(filename, message) 16 | 17 | class AverageMeter(object): 18 | """Computes and stores the average and current value""" 19 | def __init__(self): 20 | self.reset() 21 | 22 | def reset(self): 23 | self.val = 0 24 | self.avg = 0 25 | self.sum = 0 26 | self.count = 0 27 | 28 | def update(self, val, n=1): 29 | self.val = val 30 | self.sum += val * n 31 | self.count += n 32 | self.avg = self.sum / self.count 33 | 34 | def adjust_learning_rate(optimizer, epoch): 35 | """Sets the learning rate to the initial LR decayed by 10 every 3 epochs""" 36 | lr = config.lr * (0.1 ** (epoch // 3)) 37 | for param_group in optimizer.param_groups: 38 | param_group['lr'] = lr 39 | 40 | 41 | def schedule(current_epoch, current_lrs, **logs): 42 | lrs = [1e-3, 1e-4, 0.5e-4, 1e-5, 0.5e-5] 43 | epochs = [0, 1, 6, 8, 12] 44 | for lr, epoch in zip(lrs, epochs): 45 | if current_epoch >= epoch: 46 | current_lrs[5] = lr 47 | if current_epoch >= 2: 48 | current_lrs[4] = lr * 1 49 | current_lrs[3] = lr * 1 50 | current_lrs[2] = lr * 1 51 | current_lrs[1] = lr * 1 52 | current_lrs[0] = lr * 0.1 53 | return current_lrs 54 | 55 | def accuracy(output, target, topk=(1,)): 56 | """Computes the accuracy over the k top predictions for the specified values of k""" 57 | with torch.no_grad(): 58 | maxk = max(topk) 59 | batch_size = target.size(0) 60 | 61 | _, pred = output.topk(maxk, 1, True, True) 62 | pred = pred.t() 63 | correct = pred.eq(target.view(1, -1).expand_as(pred)) 64 | 65 | res = [] 66 | for k in topk: 67 | correct_k = correct[:k].view(-1).float().sum(0, keepdim=True) 68 | res.append(correct_k.mul_(100.0 / batch_size)) 69 | return res 70 | 71 | 72 | def get_learning_rate(optimizer): 73 | lr=[] 74 | for param_group in optimizer.param_groups: 75 | lr +=[ param_group['lr'] ] 76 | 77 | #assert(len(lr)==1) #we support only one param_group 78 | lr = lr[0] 79 | 80 | return lr 81 | 82 | 83 | def time_to_str(t, mode='min'): 84 | if mode=='min': 85 | t = int(t)/60 86 | hr = t//60 87 | min = t%60 88 | return '%2d hr %02d min'%(hr,min) 89 | 90 | elif mode=='sec': 91 | t = int(t) 92 | min = t//60 93 | sec = t%60 94 | return '%2d min %02d sec'%(min,sec) 95 | 96 | 97 | else: 98 | raise NotImplementedError 99 | 100 | 101 | --------------------------------------------------------------------------------