|
|
import os |
|
|
import json |
|
|
import time |
|
|
import argparse |
|
|
import pathlib |
|
|
from tqdm import tqdm |
|
|
import matplotlib.pyplot as plt |
|
|
import torch |
|
|
import torch.nn as nn |
|
|
import torch.nn.functional as F |
|
|
from torchvision import datasets |
|
|
from torch.utils.data import DataLoader |
|
|
import torchvision.transforms as transforms |
|
|
from torch.optim.lr_scheduler import _LRScheduler |
|
|
import traceback |
|
|
|
|
|
CIFAR100_TRAIN_MEAN = (0.5070751592371323, 0.48654887331495095, 0.4409178433670343) |
|
|
CIFAR100_TRAIN_STD = (0.2673342858792401, 0.2564384629170883, 0.27615047132568404) |
|
|
MILESTONES = [60, 120, 160] |
|
|
|
|
|
|
|
|
class WideBasicBlock(nn.Module): |
|
|
def __init__(self, in_planes, out_planes, dropout_rate, stride=1): |
|
|
super(WideBasicBlock, self).__init__() |
|
|
self.bn1 = nn.BatchNorm2d(in_planes) |
|
|
self.conv1 = nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False) |
|
|
self.dropout = nn.Dropout(p=dropout_rate) |
|
|
self.bn2 = nn.BatchNorm2d(out_planes) |
|
|
self.conv2 = nn.Conv2d(out_planes, out_planes, kernel_size=3, stride=1, padding=1, bias=False) |
|
|
self.relu = nn.ReLU(inplace=True) |
|
|
|
|
|
if in_planes != out_planes: |
|
|
self.shortcut = nn.Conv2d( |
|
|
in_planes, |
|
|
out_planes, |
|
|
kernel_size=1, |
|
|
stride=stride, |
|
|
padding=0, |
|
|
bias=False, |
|
|
) |
|
|
else: |
|
|
self.shortcut = nn.Identity() |
|
|
|
|
|
def forward(self, x): |
|
|
out = self.relu(self.bn1(x)) |
|
|
skip_x = x if isinstance(self.shortcut, nn.Identity) else out |
|
|
|
|
|
out = self.conv1(out) |
|
|
out = self.relu(self.bn2(out)) |
|
|
out = self.dropout(out) |
|
|
out = self.conv2(out) |
|
|
out += self.shortcut(skip_x) |
|
|
|
|
|
return out |
|
|
|
|
|
|
|
|
class WideResNet(nn.Module): |
|
|
def __init__(self, depth, widen_factor, num_classes, dropout_rate): |
|
|
super(WideResNet, self).__init__() |
|
|
|
|
|
assert (depth - 4) % 6 == 0, "Wide-resnet depth should be 6n+4" |
|
|
n = (depth - 4) / 6 |
|
|
|
|
|
n_stages = [16, 16 * widen_factor, 32 * widen_factor, 64 * widen_factor] |
|
|
|
|
|
self.conv1 = nn.Conv2d(3, n_stages[0], kernel_size=3, stride=1, padding=1, bias=False) |
|
|
self.stage1 = self._make_wide_stage(WideBasicBlock, n_stages[0], n_stages[1], n, dropout_rate, stride=1) |
|
|
self.stage2 = self._make_wide_stage(WideBasicBlock, n_stages[1], n_stages[2], n, dropout_rate, stride=2) |
|
|
self.stage3 = self._make_wide_stage(WideBasicBlock, n_stages[2], n_stages[3], n, dropout_rate, stride=2) |
|
|
self.bn1 = nn.BatchNorm2d(n_stages[3]) |
|
|
self.relu = nn.ReLU(inplace=True) |
|
|
self.avg_pool = nn.AdaptiveAvgPool2d((1, 1)) |
|
|
self.linear = nn.Linear(n_stages[3], num_classes) |
|
|
|
|
|
self._init_params() |
|
|
|
|
|
@staticmethod |
|
|
def _make_wide_stage(block, in_planes, out_planes, num_blocks, dropout_rate, stride): |
|
|
stride_list = [stride] + [1] * (int(num_blocks) - 1) |
|
|
in_planes_list = [in_planes] + [out_planes] * (int(num_blocks) - 1) |
|
|
blocks = [] |
|
|
|
|
|
for _in_planes, _stride in zip(in_planes_list, stride_list): |
|
|
blocks.append(block(_in_planes, out_planes, dropout_rate, _stride)) |
|
|
|
|
|
return nn.Sequential(*blocks) |
|
|
|
|
|
def _init_params(self): |
|
|
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): |
|
|
if m.affine: |
|
|
m.weight.data.fill_(1) |
|
|
m.bias.data.zero_() |
|
|
elif isinstance(m, nn.Linear): |
|
|
if m.bias is not None: |
|
|
m.bias.data.zero_() |
|
|
|
|
|
def forward(self, x): |
|
|
out = self.conv1(x) |
|
|
out = self.stage1(out) |
|
|
out = self.stage2(out) |
|
|
out = self.stage3(out) |
|
|
out = self.relu(self.bn1(out)) |
|
|
out = self.avg_pool(out) |
|
|
out = out.view(out.size(0), -1) |
|
|
out = self.linear(out) |
|
|
|
|
|
return out |
|
|
|
|
|
|
|
|
def wide_resnet_28_10_old(): |
|
|
return WideResNet( |
|
|
depth=28, |
|
|
widen_factor=10, |
|
|
num_classes=100, |
|
|
dropout_rate=0.0, |
|
|
) |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
parser = argparse.ArgumentParser() |
|
|
parser.add_argument("--batch_size", type=int, default=128) |
|
|
parser.add_argument("--num_workers", type=int, default=4) |
|
|
parser.add_argument("--out_dir", type=str, default="run_1") |
|
|
parser.add_argument("--in_channels", type=int, default=3) |
|
|
parser.add_argument("--data_root", type=str, default='./datasets/cifar100/') |
|
|
parser.add_argument("--learning_rate", type=float, default=0.1) |
|
|
parser.add_argument("", type=int, default=200) |
|
|
parser.add_argument("--val_per_epoch", type=int, default=5) |
|
|
config = parser.parse_args() |
|
|
|
|
|
|
|
|
try: |
|
|
final_infos = {} |
|
|
all_results = {} |
|
|
|
|
|
pathlib.Path(config.out_dir).mkdir(parents=True, exist_ok=True) |
|
|
|
|
|
model = wide_resnet_28_10_old().cuda() |
|
|
transform_train = transforms.Compose([ |
|
|
transforms.ToTensor(), |
|
|
transforms.Lambda(lambda x: F.pad(x.unsqueeze(0), |
|
|
(4, 4, 4, 4), mode='reflect').squeeze()), |
|
|
transforms.ToPILImage(), |
|
|
transforms.RandomCrop(32), |
|
|
transforms.RandomHorizontalFlip(), |
|
|
transforms.ToTensor(), |
|
|
transforms.Normalize(CIFAR100_TRAIN_MEAN, CIFAR100_TRAIN_STD), |
|
|
]) |
|
|
|
|
|
transform_test = transforms.Compose([ |
|
|
transforms.ToTensor(), |
|
|
transforms.Normalize(CIFAR100_TRAIN_MEAN, CIFAR100_TRAIN_STD) |
|
|
]) |
|
|
train_dataset = datasets.CIFAR100(root=config.data_root, train=True, |
|
|
download=True, transform=transform_train) |
|
|
test_dataset = datasets.CIFAR100(root=config.data_root, train=False, |
|
|
download=True, transform=transform_test) |
|
|
train_loader = DataLoader(train_dataset, shuffle=True, num_workers=config.num_workers, batch_size=config.batch_size) |
|
|
test_loader = DataLoader(test_dataset, shuffle=True, num_workers=config.num_workers, batch_size=config.batch_size) |
|
|
|
|
|
criterion = nn.CrossEntropyLoss().cuda() |
|
|
optimizer = torch.optim.SGD(model.parameters(), lr=config.learning_rate, momentum=0.9, weight_decay=5e-4, |
|
|
nesterov=True) |
|
|
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, len(train_loader) * config.max_epoch) |
|
|
|
|
|
best_acc = 0.0 |
|
|
start_time = time.time() |
|
|
for cur_epoch in tqdm(range(1, config.max_epoch + 1)): |
|
|
model.train() |
|
|
for batch_idx, (images, labels) in enumerate(tqdm(train_loader)): |
|
|
images, labels = images.cuda(), labels.cuda() |
|
|
optimizer.zero_grad() |
|
|
outputs = model(images) |
|
|
loss = criterion(outputs, labels) |
|
|
loss.backward() |
|
|
optimizer.step() |
|
|
scheduler.step() |
|
|
|
|
|
print(f'Finished epoch {cur_epoch} training.') |
|
|
|
|
|
if (cur_epoch % config.val_per_epoch == 0 and cur_epoch != 0) or cur_epoch == (config.max_epoch - 1): |
|
|
model.eval() |
|
|
correct = 0.0 |
|
|
for images, labels in tqdm(test_loader): |
|
|
images, labels = images.cuda(), labels.cuda() |
|
|
with torch.no_grad(): |
|
|
outputs = model(images) |
|
|
|
|
|
_, preds = outputs.max(1) |
|
|
correct += preds.eq(labels).sum() |
|
|
cur_acc = correct.float() / len(test_loader.dataset) |
|
|
print(f"Epoch: {cur_epoch}, Accuracy: {correct.float() / len(test_loader.dataset)}") |
|
|
|
|
|
if cur_acc > best_acc: |
|
|
best_acc = cur_acc |
|
|
best_epoch = cur_epoch |
|
|
torch.save(model.state_dict(), os.path.join(config.out_dir, 'best.pth')) |
|
|
|
|
|
final_infos = { |
|
|
"cifar100": { |
|
|
"means": { |
|
|
"best_acc": best_acc.item(), |
|
|
"epoch": best_epoch |
|
|
} |
|
|
} |
|
|
} |
|
|
|
|
|
with open(os.path.join(config.out_dir, "final_info.json"), "w") as f: |
|
|
json.dump(final_infos, f) |
|
|
|
|
|
except Exception as e: |
|
|
print("Original error in subprocess:", flush=True) |
|
|
traceback.print_exc(file=open(os.path.join(config.out_dir, "traceback.log"), "w")) |
|
|
raise |
