2020-1-capstone-design2 / 2016104167

Go to a project
Toggle navigation Toggle navigation pinning
조현아
6194ff3e 1 parent 6c58c61a

rm ori FAA2

Expand all
Inline Side-by-side
Showing 21 changed files with 0 additions and 1239 deletions
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
.hypothesis/
.pytest_cache/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
target/
# Jupyter Notebook
.ipynb_checkpoints
# pyenv
.python-version
# celery beat schedule file
celerybeat-schedule
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
# Fast Autoaugment
<img src="figures/faa.png" width=800px>
A Pytorch Implementation of [Fast AutoAugment](https://arxiv.org/pdf/1905.00397.pdf) and [EfficientNet](https://arxiv.org/abs/1905.11946).
## Prerequisite
* torch==1.1.0
* torchvision==0.2.2
* hyperopt==0.1.2
* future==0.17.1
* tb-nightly==1.15.0a20190622
## Usage
### Training
#### CIFAR10
```bash
# ResNet20 (w/o FastAutoAugment)
python train.py --seed=24 --scale=3 --optimizer=sgd --fast_auto_augment=False
# ResNet20 (w/ FastAutoAugment)
python train.py --seed=24 --scale=3 --optimizer=sgd --fast_auto_augment=True
# ResNet20 (w/ FastAutoAugment, Pre-found policy)
python train.py --seed=24 --scale=3 --optimizer=sgd --fast_auto_augment=True \
--augment_path=runs/ResNet_Scale3_FastAutoAugment/augmentation.cp
# ResNet32 (w/o FastAutoAugment)
python train.py --seed=24 --scale=5 --optimizer=sgd --fast_auto_augment=False
# ResNet32 (w/ FastAutoAugment)
python train.py --seed=24 --scale=5 --optimizer=sgd --fast_auto_augment=True
# EfficientNet (w/ FastAutoAugment)
python train.py --seed=24 --pi=0 --optimizer=adam --fast_auto_augment=True \
--network=efficientnet_cifar10 --activation=swish
```
#### ImageNet (You can use any backbone networks in [torchvision.models](https://pytorch.org/docs/stable/torchvision/models.html))
```bash
# BaseNet (w/o FastAutoAugment)
python train.py --seed=24 --dataset=imagenet --optimizer=adam --network=resnet50
# EfficientNet (w/ FastAutoAugment) (UnderConstruction)
python train.py --seed=24 --dataset=imagenet --pi=0 --optimizer=adam --fast_auto_augment=True \
--network=efficientnet --activation=swish
```
### Eval
```bash
# Single Image testing
python eval.py --model_path=runs/ResNet_Scale3_Basline
# 5-crops testing
python eval.py --model_path=runs/ResNet_Scale3_Basline --five_crops=True
```
## Experiments
### Fast AutoAugment
#### ResNet20 (CIFAR10)
* Pre-trained model [[Download](https://drive.google.com/file/d/12D8050yGGiKWGt8_R8QTlkoQ6wq_icBn/view?usp=sharing)]
* Validation Curve
<img src="figures/resnet20_valid.png">
* Evaluation (Acc @1)
| | Valid | Test(Single) |
|----------------|-------|-------------|
| ResNet20 | 90.70 | **91.45** |
| ResNet20 + FAA |**92.46**| **91.45** |
#### ResNet34 (CIFAR10)
* Validation Curve
<img src="figures/resnet34_valid.png">
* Evaluation (Acc @1)
| | Valid | Test(Single) |
|----------------|-------|-------------|
| ResNet34 | 91.54 | 91.47 |
| ResNet34 + FAA |**92.76**| **91.99** |
### Found Policy [[Download](https://drive.google.com/file/d/1Ia_IxPY3-T7m8biyl3QpxV1s5EA5gRDF/view?usp=sharing)]
<img src="figures/pm.png">
### Augmented images
<img src="figures/augmented_images.png">
<img src="figures/augmented_images2.png">
import os
import fire
import json
from pprint import pprint
import torch
import torch.nn as nn
from utils import *
def eval(model_path):
print('\n[+] Parse arguments')
kwargs_path = os.path.join(model_path, 'kwargs.json')
kwargs = json.loads(open(kwargs_path).read())
args, kwargs = parse_args(kwargs)
pprint(args)
device = torch.device('cuda' if args.use_cuda else 'cpu')
print('\n[+] Create network')
model = select_model(args)
optimizer = select_optimizer(args, model)
criterion = nn.CrossEntropyLoss()
if args.use_cuda:
model = model.cuda()
criterion = criterion.cuda()
print('\n[+] Load model')
weight_path = os.path.join(model_path, 'model', 'model.pt')
model.load_state_dict(torch.load(weight_path))
print('\n[+] Load dataset')
test_transform = get_valid_transform(args, model)
test_dataset = get_dataset(args, test_transform, 'test')
test_loader = iter(get_dataloader(args, test_dataset))
print('\n[+] Start testing')
_test_res = validate(args, model, criterion, test_loader, step=0, writer=None)
print('\n[+] Valid results')
print(' Acc@1 : {:.3f}%'.format(_test_res[0].data.cpu().numpy()[0]*100))
print(' Acc@5 : {:.3f}%'.format(_test_res[1].data.cpu().numpy()[0]*100))
print(' Loss : {:.3f}'.format(_test_res[2].data))
print(' Infer Time(per image) : {:.3f}ms'.format(_test_res[3]*1000 / len(test_dataset)))
if __name__ == '__main__':
fire.Fire(eval)
import copy
import json
import time
import torch
import random
import torchvision.transforms as transforms
from torch.utils.data import Subset
from sklearn.model_selection import StratifiedShuffleSplit
from concurrent.futures import ProcessPoolExecutor
from transforms import *
from hyperopt import fmin, tpe, hp, STATUS_OK, Trials
from utils import *
DEFALUT_CANDIDATES = [
ShearXY,
TranslateXY,
Rotate,
AutoContrast,
Invert,
Equalize,
Solarize,
Posterize,
Contrast,
Color,
Brightness,
Sharpness,
Cutout,
# SamplePairing,
]
def train_child(args, model, dataset, subset_indx, device=None):
optimizer = select_optimizer(args, model)
scheduler = select_scheduler(args, optimizer)
criterion = nn.CrossEntropyLoss()
dataset.transform = transforms.Compose([
transforms.Resize(32),
transforms.ToTensor()])
subset = Subset(dataset, subset_indx)
data_loader = get_inf_dataloader(args, subset)
if device:
model = model.to(device)
criterion = criterion.to(device)
elif args.use_cuda:
model = model.cuda()
criterion = criterion.cuda()
if torch.cuda.device_count() > 1:
print('\n[+] Use {} GPUs'.format(torch.cuda.device_count()))
model = nn.DataParallel(model)
start_t = time.time()
for step in range(args.start_step, args.max_step):
batch = next(data_loader)
_train_res = train_step(args, model, optimizer, scheduler, criterion, batch, step, None, device)
if step % args.print_step == 0:
print('\n[+] Training step: {}/{}\tElapsed time: {:.2f}min\tLearning rate: {}\tDevice: {}'.format(
step, args.max_step,(time.time()-start_t)/60, optimizer.param_groups[0]['lr'], device))
print(' Acc@1 : {:.3f}%'.format(_train_res[0].data.cpu().numpy()[0]*100))
print(' Acc@5 : {:.3f}%'.format(_train_res[1].data.cpu().numpy()[0]*100))
print(' Loss : {}'.format(_train_res[2].data))
return _train_res
def validate_child(args, model, dataset, subset_indx, transform, device=None):
criterion = nn.CrossEntropyLoss()
if device:
model = model.to(device)
criterion = criterion.to(device)
elif args.use_cuda:
model = model.cuda()
criterion = criterion.cuda()
dataset.transform = transform
subset = Subset(dataset, subset_indx)
data_loader = get_dataloader(args, subset, pin_memory=False)
return validate(args, model, criterion, data_loader, 0, None, device)
def get_next_subpolicy(transform_candidates, op_per_subpolicy=2):
n_candidates = len(transform_candidates)
subpolicy = []
for i in range(op_per_subpolicy):
indx = random.randrange(n_candidates)
prob = random.random()
mag = random.random()
subpolicy.append(transform_candidates[indx](prob, mag))
subpolicy = transforms.Compose([
*subpolicy,
transforms.Resize(32),
transforms.ToTensor()])
return subpolicy
def search_subpolicies(args, transform_candidates, child_model, dataset, Da_indx, B, device):
subpolicies = []
for b in range(B):
subpolicy = get_next_subpolicy(transform_candidates)
val_res = validate_child(args, child_model, dataset, Da_indx, subpolicy, device)
subpolicies.append((subpolicy, val_res[2]))
return subpolicies
def search_subpolicies_hyperopt(args, transform_candidates, child_model, dataset, Da_indx, B, device):
def _objective(sampled):
subpolicy = [transform(prob, mag)
for transform, prob, mag in sampled]
subpolicy = transforms.Compose([
transforms.Resize(32),
*subpolicy,
transforms.ToTensor()])
val_res = validate_child(args, child_model, dataset, Da_indx, subpolicy, device)
loss = val_res[2].cpu().numpy()
return {'loss': loss, 'status': STATUS_OK }
space = [(hp.choice('transform1', transform_candidates), hp.uniform('prob1', 0, 1.0), hp.uniform('mag1', 0, 1.0)),
(hp.choice('transform2', transform_candidates), hp.uniform('prob2', 0, 1.0), hp.uniform('mag2', 0, 1.0))]
trials = Trials()
best = fmin(_objective,
space=space,
algo=tpe.suggest,
max_evals=B,
trials=trials)
subpolicies = []
for t in trials.trials:
vals = t['misc']['vals']
subpolicy = [transform_candidates[vals['transform1'][0]](vals['prob1'][0], vals['mag1'][0]),
transform_candidates[vals['transform2'][0]](vals['prob2'][0], vals['mag2'][0])]
subpolicy = transforms.Compose([
## baseline augmentation
transforms.Pad(4),
transforms.RandomCrop(32),
transforms.RandomHorizontalFlip(),
## policy
*subpolicy,
## to tensor
transforms.ToTensor()])
subpolicies.append((subpolicy, t['result']['loss']))
return subpolicies
def get_topn_subpolicies(subpolicies, N=10):
return sorted(subpolicies, key=lambda subpolicy: subpolicy[1])[:N]
def process_fn(args_str, model, dataset, Dm_indx, Da_indx, T, transform_candidates, B, N, k):
kwargs = json.loads(args_str)
args, kwargs = parse_args(kwargs)
device_id = k % torch.cuda.device_count()
device = torch.device('cuda:%d' % device_id)
_transform = []
print('[+] Child %d training strated (GPU: %d)' % (k, device_id))
# train child model
child_model = copy.deepcopy(model)
train_res = train_child(args, child_model, dataset, Dm_indx, device)
# search sub policy
for t in range(T):
#subpolicies = search_subpolicies(args, transform_candidates, child_model, dataset, Da_indx, B, device)
subpolicies = search_subpolicies_hyperopt(args, transform_candidates, child_model, dataset, Da_indx, B, device)
subpolicies = get_topn_subpolicies(subpolicies, N)
_transform.extend([subpolicy[0] for subpolicy in subpolicies])
return _transform
def fast_auto_augment(args, model, transform_candidates=None, K=5, B=100, T=2, N=10, num_process=5):
args_str = json.dumps(args._asdict())
dataset = get_dataset(args, None, 'trainval')
num_process = min(torch.cuda.device_count(), num_process)
transform, futures = [], []
torch.multiprocessing.set_start_method('spawn', force=True)
if not transform_candidates:
transform_candidates = DEFALUT_CANDIDATES
# split
Dm_indexes, Da_indexes = split_dataset(args, dataset, K)
with ProcessPoolExecutor(max_workers=num_process) as executor:
for k, (Dm_indx, Da_indx) in enumerate(zip(Dm_indexes, Da_indexes)):
future = executor.submit(process_fn,
args_str, model, dataset, Dm_indx, Da_indx, T, transform_candidates, B, N, k)
futures.append(future)
for future in futures:
transform.extend(future.result())
transform = transforms.RandomChoice(transform)
return transform
import torch.nn as nn
class BaseNet(nn.Module):
def __init__(self, backbone, args):
super(BaseNet, self).__init__()
# Separate layers
self.first = nn.Sequential(*list(backbone.children())[:1])
self.after = nn.Sequential(*list(backbone.children())[1:-1])
self.fc = list(backbone.children())[-1]
self.img_size = (224, 224)
def forward(self, x):
f = self.first(x)
x = self.after(f)
x = x.reshape(x.size(0), -1)
x = self.fc(x)
return x, f
import math
import torch.nn as nn
import torch.nn.functional as F
def round_fn(orig, multiplier):
if not multiplier:
return orig
return int(math.ceil(multiplier * orig))
def get_activation_fn(activation):
if activation == "swish":
return Swish
elif activation == "relu":
return nn.ReLU
else:
raise Exception('Unkown activation %s' % activation)
class Swish(nn.Module):
""" Swish activation function, s(x) = x * sigmoid(x) """
def __init__(self, inplace=False):
super().__init__()
self.inplace = True
def forward(self, x):
if self.inplace:
x.mul_(F.sigmoid(x))
return x
else:
return x * F.sigmoid(x)
class ConvBlock(nn.Module):
""" Conv + BatchNorm + Activation """
def __init__(self, in_channel, out_channel, kernel_size,
padding=0, stride=1, activation="swish"):
super().__init__()
self.fw = nn.Sequential(
nn.Conv2d(in_channel, out_channel, kernel_size,
padding=padding, stride=stride, bias=False),
nn.BatchNorm2d(out_channel),
get_activation_fn(activation)())
def forward(self, x):
return self.fw(x)
class DepthwiseConvBlock(nn.Module):
""" DepthwiseConv2D + BatchNorm + Activation """
def __init__(self, in_channel, kernel_size,
padding=0, stride=1, activation="swish"):
super().__init__()
self.fw = nn.Sequential(
nn.Conv2d(in_channel, in_channel, kernel_size,
padding=padding, stride=stride, groups=in_channel, bias=False),
nn.BatchNorm2d(in_channel),
get_activation_fn(activation)())
def forward(self, x):
return self.fw(x)
class MBConv(nn.Module):
""" Inverted residual block """
def __init__(self, in_channel, out_channel, kernel_size,
stride=1, expand_ratio=1, activation="swish"):
super().__init__()
self.in_channel = in_channel
self.out_channel = out_channel
self.expand_ratio = expand_ratio
self.stride = stride
if expand_ratio != 1:
self.expand = ConvBlock(in_channel, in_channel*expand_ratio, 1,
activation=activation)
self.dw_conv = DepthwiseConvBlock(in_channel*expand_ratio, kernel_size,
padding=(kernel_size-1)//2,
stride=stride, activation=activation)
self.pw_conv = ConvBlock(in_channel*expand_ratio, out_channel, 1,
activation=activation)
def forward(self, inputs):
if self.expand_ratio != 1:
x = self.expand(inputs)
else:
x = inputs
x = self.dw_conv(x)
x = self.pw_conv(x)
if self.in_channel == self.out_channel and \
self.stride == 1:
x = x + inputs
return x
class Net(nn.Module):
""" EfficientNet """
def __init__(self, args):
super(Net, self).__init__()
pi = args.pi
activation = args.activation
num_classes = args.num_classes
self.d = 1.2 ** pi
self.w = 1.1 ** pi
self.r = 1.15 ** pi
self.img_size = (round_fn(224, self.r), round_fn(224, self.r))
self.stage1 = ConvBlock(3, round_fn(32, self.w),
kernel_size=3, padding=1, stride=2, activation=activation)
self.stage2 = self.make_layers(round_fn(32, self.w), round_fn(16, self.w),
depth=round_fn(1, self.d), kernel_size=3,
half_resolution=False, expand_ratio=1, activation=activation)
self.stage3 = self.make_layers(round_fn(16, self.w), round_fn(24, self.w),
depth=round_fn(2, self.d), kernel_size=3,
half_resolution=True, expand_ratio=6, activation=activation)
self.stage4 = self.make_layers(round_fn(24, self.w), round_fn(40, self.w),
depth=round_fn(2, self.d), kernel_size=5,
half_resolution=True, expand_ratio=6, activation=activation)
self.stage5 = self.make_layers(round_fn(40, self.w), round_fn(80, self.w),
depth=round_fn(3, self.d), kernel_size=3,
half_resolution=True, expand_ratio=6, activation=activation)
self.stage6 = self.make_layers(round_fn(80, self.w), round_fn(112, self.w),
depth=round_fn(3, self.d), kernel_size=5,
half_resolution=False, expand_ratio=6, activation=activation)
self.stage7 = self.make_layers(round_fn(112, self.w), round_fn(192, self.w),
depth=round_fn(4, self.d), kernel_size=5,
half_resolution=True, expand_ratio=6, activation=activation)
self.stage8 = self.make_layers(round_fn(192, self.w), round_fn(320, self.w),
depth=round_fn(1, self.d), kernel_size=3,
half_resolution=False, expand_ratio=6, activation=activation)
self.stage9 = ConvBlock(round_fn(320, self.w), round_fn(1280, self.w),
kernel_size=1, activation=activation)
self.fc = nn.Linear(round_fn(7*7*1280, self.w), num_classes)
def make_layers(self, in_channel, out_channel, depth, kernel_size,
half_resolution=False, expand_ratio=1, activation="swish"):
blocks = []
for i in range(depth):
stride = 2 if half_resolution and i==0 else 1
blocks.append(
MBConv(in_channel, out_channel, kernel_size,
stride=stride, expand_ratio=expand_ratio, activation=activation))
in_channel = out_channel
return nn.Sequential(*blocks)
def forward(self, x):
assert x.size()[-2:] == self.img_size, \
'Image size must be %r, but %r given' % (self.img_size, x.size()[-2])
x = self.stage1(x)
x = self.stage2(x)
x = self.stage3(x)
x = self.stage4(x)
x = self.stage5(x)
x = self.stage6(x)
x = self.stage7(x)
x = self.stage8(x)
x = self.stage9(x)
x = x.reshape(x.size(0), -1)
x = self.fc(x)
return x, x
import math
import torch.nn as nn
import torch.nn.functional as F
def round_fn(orig, multiplier):
if not multiplier:
return orig
return int(math.ceil(multiplier * orig))
def get_activation_fn(activation):
if activation == "swish":
return Swish
elif activation == "relu":
return nn.ReLU
else:
raise Exception('Unkown activation %s' % activation)
class Swish(nn.Module):
""" Swish activation function, s(x) = x * sigmoid(x) """
def __init__(self, inplace=False):
super().__init__()
self.inplace = True
def forward(self, x):
if self.inplace:
x.mul_(F.sigmoid(x))
return x
else:
return x * F.sigmoid(x)
class ConvBlock(nn.Module):
""" Conv + BatchNorm + Activation """
def __init__(self, in_channel, out_channel, kernel_size,
padding=0, stride=1, activation="swish"):
super().__init__()
self.fw = nn.Sequential(
nn.Conv2d(in_channel, out_channel, kernel_size,
padding=padding, stride=stride, bias=False),
nn.BatchNorm2d(out_channel),
get_activation_fn(activation)())
def forward(self, x):
return self.fw(x)
class DepthwiseConvBlock(nn.Module):
""" DepthwiseConv2D + BatchNorm + Activation """
def __init__(self, in_channel, kernel_size,
padding=0, stride=1, activation="swish"):
super().__init__()
self.fw = nn.Sequential(
nn.Conv2d(in_channel, in_channel, kernel_size,
padding=padding, stride=stride, groups=in_channel, bias=False),
nn.BatchNorm2d(in_channel),
get_activation_fn(activation)())
def forward(self, x):
return self.fw(x)
class SEBlock(nn.Module):
""" Squeeze and Excitation Block """
def __init__(self, in_channel, se_ratio=16):
super().__init__()
self.global_avgpool = nn.AdaptiveAvgPool2d((1,1))
inter_channel = in_channel // se_ratio
self.reduce = nn.Sequential(
nn.Conv2d(in_channel, inter_channel,
kernel_size=1, padding=0, stride=1),
nn.ReLU())
self.expand = nn.Sequential(
nn.Conv2d(inter_channel, in_channel,
kernel_size=1, padding=0, stride=1),
nn.Sigmoid())
def forward(self, x):
s = self.global_avgpool(x)
s = self.reduce(s)
s = self.expand(s)
return x * s
class MBConv(nn.Module):
""" Inverted residual block """
def __init__(self, in_channel, out_channel, kernel_size,
stride=1, expand_ratio=1, activation="swish", use_seblock=False):
super().__init__()
self.in_channel = in_channel
self.out_channel = out_channel
self.expand_ratio = expand_ratio
self.stride = stride
self.use_seblock = use_seblock
if expand_ratio != 1:
self.expand = ConvBlock(in_channel, in_channel*expand_ratio, 1,
activation=activation)
self.dw_conv = DepthwiseConvBlock(in_channel*expand_ratio, kernel_size,
padding=(kernel_size-1)//2,
stride=stride, activation=activation)
if use_seblock:
self.seblock = SEBlock(in_channel*expand_ratio)
self.pw_conv = ConvBlock(in_channel*expand_ratio, out_channel, 1,
activation=activation)
def forward(self, inputs):
if self.expand_ratio != 1:
x = self.expand(inputs)
else:
x = inputs
x = self.dw_conv(x)
if self.use_seblock:
x = self.seblock(x)
x = self.pw_conv(x)
if self.in_channel == self.out_channel and \
self.stride == 1:
x = x + inputs
return x
class Net(nn.Module):
""" EfficientNet """
def __init__(self, args):
super(Net, self).__init__()
pi = args.pi
activation = args.activation
num_classes = 10
self.d = 1.2 ** pi
self.w = 1.1 ** pi
self.r = 1.15 ** pi
self.img_size = (round_fn(32, self.r), round_fn(32, self.r))
self.use_seblock = args.use_seblock
self.stage1 = ConvBlock(3, round_fn(32, self.w),
kernel_size=3, padding=1, stride=2, activation=activation)
self.stage2 = self.make_layers(round_fn(32, self.w), round_fn(16, self.w),
depth=round_fn(1, self.d), kernel_size=3,
half_resolution=False, expand_ratio=1, activation=activation)
self.stage3 = self.make_layers(round_fn(16, self.w), round_fn(24, self.w),
depth=round_fn(2, self.d), kernel_size=3,
half_resolution=True, expand_ratio=6, activation=activation)
self.stage4 = self.make_layers(round_fn(24, self.w), round_fn(40, self.w),
depth=round_fn(2, self.d), kernel_size=5,
half_resolution=True, expand_ratio=6, activation=activation)
self.stage5 = self.make_layers(round_fn(40, self.w), round_fn(80, self.w),
depth=round_fn(3, self.d), kernel_size=3,
half_resolution=True, expand_ratio=6, activation=activation)
self.stage6 = self.make_layers(round_fn(80, self.w), round_fn(112, self.w),
depth=round_fn(3, self.d), kernel_size=5,
half_resolution=False, expand_ratio=6, activation=activation)
self.stage7 = self.make_layers(round_fn(112, self.w), round_fn(192, self.w),
depth=round_fn(4, self.d), kernel_size=5,
half_resolution=True, expand_ratio=6, activation=activation)
self.stage8 = self.make_layers(round_fn(192, self.w), round_fn(320, self.w),
depth=round_fn(1, self.d), kernel_size=3,
half_resolution=False, expand_ratio=6, activation=activation)
self.stage9 = ConvBlock(round_fn(320, self.w), round_fn(1280, self.w),
kernel_size=1, activation=activation)
self.fc = nn.Linear(round_fn(1280, self.w), num_classes)
def make_layers(self, in_channel, out_channel, depth, kernel_size,
half_resolution=False, expand_ratio=1, activation="swish"):
blocks = []
for i in range(depth):
stride = 2 if half_resolution and i==0 else 1
blocks.append(
MBConv(in_channel, out_channel, kernel_size,
stride=stride, expand_ratio=expand_ratio, activation=activation, use_seblock=self.use_seblock))
in_channel = out_channel
return nn.Sequential(*blocks)
def forward(self, x):
assert x.size()[-2:] == self.img_size, \
'Image size must be %r, but %r given' % (self.img_size, x.size()[-2])
s = self.stage1(x)
x = self.stage2(s)
x = self.stage3(x)
x = self.stage4(x)
x = self.stage5(x)
x = self.stage6(x)
x = self.stage7(x)
x = self.stage8(x)
x = self.stage9(x)
x = x.reshape(x.size(0), -1)
x = self.fc(x)
return x, s
import torch.nn as nn
class ResidualBlock(nn.Module):
def __init__(self, in_channel, out_channel, stride):
super(ResidualBlock, self).__init__()
self.in_channel = in_channel
self.out_channel = out_channel
self.stride = stride
self.conv1 = nn.Sequential(
nn.Conv2d(in_channel, out_channel,
kernel_size=3, padding=1, stride=stride),
nn.BatchNorm2d(out_channel))
self.relu = nn.ReLU(inplace=True)
self.conv2 = nn.Sequential(
nn.Conv2d(out_channel, out_channel,
kernel_size=3, padding=1),
nn.BatchNorm2d(out_channel))
if self.in_channel != self.out_channel or \
self.stride != 1:
self.down = nn.Sequential(
nn.Conv2d(in_channel, out_channel,
kernel_size=1, stride=stride),
nn.BatchNorm2d(out_channel))
def forward(self, b):
t = self.conv1(b)
t = self.relu(t)
t = self.conv2(t)
if self.in_channel != self.out_channel or \
self.stride != 1:
b = self.down(b)
t += b
t = self.relu(t)
return t
class Net(nn.Module):
def __init__(self, args):
super(Net, self).__init__()
scale = args.scale
self.stem = nn.Sequential(
nn.Conv2d(3, 16,
kernel_size=3, padding=1),
nn.BatchNorm2d(16),
nn.ReLU(inplace=True))
self.layer1 = nn.Sequential(*[
ResidualBlock(16, 16, 1) for _ in range(2*scale)])
self.layer2 = nn.Sequential(*[
ResidualBlock(in_channel=(16 if i==0 else 32),
out_channel=32,
stride=(2 if i==0 else 1)) for i in range(2*scale)])
self.layer3 = nn.Sequential(*[
ResidualBlock(in_channel=(32 if i==0 else 64),
out_channel=64,
stride=(2 if i==0 else 1)) for i in range(2*scale)])
self.avg_pool = nn.AdaptiveAvgPool2d((1,1))
self.fc = nn.Linear(64, 10)
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
def forward(self, x):
s = self.stem(x)
x = self.layer1(s)
x = self.layer2(x)
x = self.layer3(x)
x = self.avg_pool(x)
x = x.reshape(x.size(0), -1)
x = self.fc(x)
return x, s
import os
import fire
import time
import json
import random
from pprint import pprint
import torch.nn as nn
import torch.backends.cudnn as cudnn
from torch.utils.tensorboard import SummaryWriter
from networks import *
from utils import *
def train(**kwargs):
print('\n[+] Parse arguments')
args, kwargs = parse_args(kwargs)
pprint(args)
device = torch.device('cuda' if args.use_cuda else 'cpu')
print('\n[+] Create log dir')
model_name = get_model_name(args)
log_dir = os.path.join('./runs', model_name)
os.makedirs(os.path.join(log_dir, 'model'))
json.dump(kwargs, open(os.path.join(log_dir, 'kwargs.json'), 'w'))
writer = SummaryWriter(log_dir=log_dir)
if args.seed is not None:
random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.deterministic = True
print('\n[+] Create network')
model = select_model(args)
optimizer = select_optimizer(args, model)
scheduler = select_scheduler(args, optimizer)
criterion = nn.CrossEntropyLoss()
if args.use_cuda:
model = model.cuda()
criterion = criterion.cuda()
#writer.add_graph(model)
print('\n[+] Load dataset')
transform = get_train_transform(args, model, log_dir)
val_transform = get_valid_transform(args, model)
train_dataset = get_dataset(args, transform, 'train')
valid_dataset = get_dataset(args, val_transform, 'val')
train_loader = iter(get_inf_dataloader(args, train_dataset))
max_epoch = len(train_dataset) // args.batch_size
best_acc = -1
print('\n[+] Start training')
if torch.cuda.device_count() > 1:
print('\n[+] Use {} GPUs'.format(torch.cuda.device_count()))
model = nn.DataParallel(model)
start_t = time.time()
for step in range(args.start_step, args.max_step):
batch = next(train_loader)
_train_res = train_step(args, model, optimizer, scheduler, criterion, batch, step, writer)
if step % args.print_step == 0:
print('\n[+] Training step: {}/{}\tTraining epoch: {}/{}\tElapsed time: {:.2f}min\tLearning rate: {}'.format(
step, args.max_step, current_epoch, max_epoch, (time.time()-start_t)/60, optimizer.param_groups[0]['lr']))
writer.add_scalar('train/learning_rate', optimizer.param_groups[0]['lr'], global_step=step)
writer.add_scalar('train/acc1', _train_res[0], global_step=step)
writer.add_scalar('train/acc5', _train_res[1], global_step=step)
writer.add_scalar('train/loss', _train_res[2], global_step=step)
writer.add_scalar('train/forward_time', _train_res[3], global_step=step)
writer.add_scalar('train/backward_time', _train_res[4], global_step=step)
print(' Acc@1 : {:.3f}%'.format(_train_res[0].data.cpu().numpy()[0]*100))
print(' Acc@5 : {:.3f}%'.format(_train_res[1].data.cpu().numpy()[0]*100))
print(' Loss : {}'.format(_train_res[2].data))
print(' FW Time : {:.3f}ms'.format(_train_res[3]*1000))
print(' BW Time : {:.3f}ms'.format(_train_res[4]*1000))
if step % args.val_step == args.val_step-1:
valid_loader = iter(get_dataloader(args, valid_dataset))
_valid_res = validate(args, model, criterion, valid_loader, step, writer)
print('\n[+] Valid results')
writer.add_scalar('valid/acc1', _valid_res[0], global_step=step)
writer.add_scalar('valid/acc5', _valid_res[1], global_step=step)
writer.add_scalar('valid/loss', _valid_res[2], global_step=step)
print(' Acc@1 : {:.3f}%'.format(_valid_res[0].data.cpu().numpy()[0]*100))
print(' Acc@5 : {:.3f}%'.format(_valid_res[1].data.cpu().numpy()[0]*100))
print(' Loss : {}'.format(_valid_res[2].data))
if _valid_res[0] > best_acc:
best_acc = _valid_res[0]
torch.save(model.state_dict(), os.path.join(log_dir, "model","model.pt"))
print('\n[+] Model saved')
writer.close()
if __name__ == '__main__':
fire.Fire(train)
import numpy as np
import torch.nn as nn
import torchvision.transforms as transforms
from abc import ABC, abstractmethod
from PIL import Image, ImageOps, ImageEnhance
class BaseTransform(ABC):
def __init__(self, prob, mag):
self.prob = prob
self.mag = mag
def __call__(self, img):
return transforms.RandomApply([self.transform], self.prob)(img)
def __repr__(self):
return '%s(prob=%.2f, magnitude=%.2f)' % \
(self.__class__.__name__, self.prob, self.mag)
@abstractmethod
def transform(self, img):
pass
class ShearXY(BaseTransform):
def transform(self, img):
degrees = self.mag * 360
t = transforms.RandomAffine(0, shear=degrees, resample=Image.BILINEAR)
return t(img)
class TranslateXY(BaseTransform):
def transform(self, img):
translate = (self.mag, self.mag)
t = transforms.RandomAffine(0, translate=translate, resample=Image.BILINEAR)
return t(img)
class Rotate(BaseTransform):
def transform(self, img):
degrees = self.mag * 360
t = transforms.RandomRotation(degrees, Image.BILINEAR)
return t(img)
class AutoContrast(BaseTransform):
def transform(self, img):
cutoff = int(self.mag * 49)
return ImageOps.autocontrast(img, cutoff=cutoff)
class Invert(BaseTransform):
def transform(self, img):
return ImageOps.invert(img)
class Equalize(BaseTransform):
def transform(self, img):
return ImageOps.equalize(img)
class Solarize(BaseTransform):
def transform(self, img):
threshold = (1-self.mag) * 255
return ImageOps.solarize(img, threshold)
class Posterize(BaseTransform):
def transform(self, img):
bits = int((1-self.mag) * 8)
return ImageOps.posterize(img, bits=bits)
class Contrast(BaseTransform):
def transform(self, img):
factor = self.mag * 10
return ImageEnhance.Contrast(img).enhance(factor)
class Color(BaseTransform):
def transform(self, img):
factor = self.mag * 10
return ImageEnhance.Color(img).enhance(factor)
class Brightness(BaseTransform):
def transform(self, img):
factor = self.mag * 10
return ImageEnhance.Brightness(img).enhance(factor)
class Sharpness(BaseTransform):
def transform(self, img):
factor = self.mag * 10
return ImageEnhance.Sharpness(img).enhance(factor)
class Cutout(BaseTransform):
def transform(self, img):
n_holes = 1
length = 24 * self.mag
cutout_op = CutoutOp(n_holes=n_holes, length=length)
return cutout_op(img)
class CutoutOp(object):
"""
https://github.com/uoguelph-mlrg/Cutout
Randomly mask out one or more patches from an image.
Args:
n_holes (int): Number of patches to cut out of each image.
length (int): The length (in pixels) of each square patch.
"""
def __init__(self, n_holes, length):
self.n_holes = n_holes
self.length = length
def __call__(self, img):
"""
Args:
img (Tensor): Tensor image of size (C, H, W).
Returns:
Tensor: Image with n_holes of dimension length x length cut out of it.
"""
w, h = img.size
mask = np.ones((h, w, 1), np.uint8)
for n in range(self.n_holes):
y = np.random.randint(h)
x = np.random.randint(w)
y1 = np.clip(y - self.length // 2, 0, h).astype(int)
y2 = np.clip(y + self.length // 2, 0, h).astype(int)
x1 = np.clip(x - self.length // 2, 0, w).astype(int)
x2 = np.clip(x + self.length // 2, 0, w).astype(int)
mask[y1: y2, x1: x2, :] = 0.
img = mask*np.asarray(img).astype(np.uint8)
img = Image.fromarray(mask*np.asarray(img))
return img