이 글은 MixMatch: A Holistic Approach to Semi-Supervised Learning논문을 참고하여 정리하였음을 먼저 밝힙니다. 논문에서 사용한 방법론을 간단하게 설명하고, pytorch를 이용하여 코드를 구현한 후 추가적으로 설명드리겠습니다. 혹시 제가 잘못 알고 있는 점이나 보안할 점이 있다면 댓글 부탁드립니다.
0. Summary
SSL(Semi Supervised Learning)은 기존 supervised learning에 unlabeled data를 이용하여 보다 분류 성능을 향상시키고 모델을 robust하게 만드는 기법입니다. SSL 관련 연구에서는 Loss function에서 다음의 세 가지의 문제를 항상 고려합니다.
1. Entropy minimization
: 모델이 Unlabeled data에 대해 신뢰할 수 있도록 output을 예측합니다.
Semi-supervised classification은 밀도가 높은 지역을 통과하는 decision boundary가 형성되는 것을 막아야합니다. 따라서 unlabeled data에 낮은 entropy를 부여하여 그 문제를 해결합니다. 본 논문에서는 해당 부분을 2번에서 설명할 VAT에 결합하여 사용합니다.
2. Consistency regularization
: 모델이 노이즈에 robust(조금 바뀌더라도 같은 분포를 출력)한 output을 예측합니다.
따라서 simple한 경우를 생각하면, unlabeled data에 대하여 augmentation한 결과의 class 예측 분포와 augmentation을 적용하지 않은 class 예측 분포의 결과가 유사해야합니다.
즉, 위와 같은 loss함수를 주로 loss term에 추가하는 방식을 사용합니다.
그러나 augmentation은 data의 domain specific한 전략이므로 모든 데이터에 대해 각자 적합한 augmentation을 적용해야 한다는 단점이 있습니다. 따라서 본 논문에서는 VAT(Virtual Adversarial Training)이라는 방법론을 사용하였습니다. VAT는 출력 클래스 분포가 최대로 변경될 수 있는 지점까지 계산하는 additive pertubation을 사용합니다. 또한 일반적으로 주로 사용하는 random horizontal flips, crops를 사용하였습니다.
3. Generic regularization
: 모델 학습 시 training data에 overfitting되는걸 방지합니다.
본 논문에서는 해당 부분을 완화하기위해 우선 많이 사용하는 L2 norm을 모델에 추가하였습니다. 또한 SSL에서 많이 사용하는 MixUp 방법론을 labeled와 unlabeled data에 모두 적용하는 방법론을 제시합니다.
따라서 저자는 위의 세 가지 문제를 커버할 수 있는 loss fuction을 소개합니다. 그 결과 CIFAR-10 dataset에서 SOTA의 성능을 보였습니다.
1. MixMatch
MixMatch는 배치단위로 들어온 데이터에 대해 항상 augmentation한 결과를 적용합니다. Labeled data의 배치사이즈 만큼인 \(\chi\)와 동일한 배치사이즈 만큼의 unlabeled data \({\mathcal{U}}\)를 각각 augmentation 시킨 \(\chi\)'과 \({\mathcal{U}}\)'을 loss function에 사용합니다.
위의 수식에서 H(p, q)는 p와 q 분포에대한 cross-entropy function입니다. 본 논문의 loss function 또한 일반적인 semi-supervised learning처럼 labeled data와 unlabeled data의 적절한 합으로 구성되어 있습니다. 수식에 대한 자세한 내용은 아래에서 설명하겠습니다.
2. Data Augmentation
본 논문에서는 Augmentation으로 얻은 data와 unlabeled로 부터 추정된 label에 대한 notation을 다음과 같이 정의합니다.
또한 unlabeled \({u_b}\) data에 대해 추정한 label \({q_b}\)라 정의합니다.
3. Label Guessing
우선 labeled data로 학습한 모델에 K개의 augmentation 기법을 적용한 unlabeled data를 집어넣어 각 배치단위 별 각 class에 속할 확률 분포를 뽑아 평균을 취한 결과를 얻습니다.
그 결과에 sharpening이라는 기법을 적용하여 각 데이터 별 하나의 클래스가 앞도적으로 높은 확률을 갖도록 설정합니다.
T는 sharpening에 사용하는 하이퍼 파라미터이며(T->0 일때, Dirac 분포를 따름), 위에서 정의한 notation을 사용하여 나타내면 \({q_b}\) = Sharpen(\({\bar{q_b}}\),T)로 표현할 수 있습니다. 저자는 위와 같은 방법론이 모델의 lower-entropy prediction을 할 수 있다고 말합니다.
4. MixUp
본 연구의 핵심인 MixUp은 다음과 같이 \(\alpha\)라는 하이퍼 파라미터로 부터 얻은 \(\lambda\)를 사용합니다. 우선 위에서 augmentation하여 얻은 labeled data와 unlabeled data에 대한 확률분포를 얻습니다.
그 다음 얻은 결과를 concat하여 섞습니다. 논문에서는 해당 결과를 \({\mathcal{W}}\)로 정의합니다.
그 후 pair한 두 개의 데이터를 뽑아 다음과 같은 작업을 진행합니다.
\(\lambda\)'은 max 함수에 의해 항상 0.5보다 크거나 같은 값으로서 앞에 존재하는 \({x_1}\)에 더 많은 가중치를 주어 새롭게 x' 데이터와 p' class 확률분포를 만듭니다.
최종적으로 그렇게 얻은 데이터로부터 상위분포에 대하여 배치 사이즈만큼 labeled data로 분류하고 나머지를 unlabeled data로 분류합니다(학습 시 Unlabeled와 labeled에서 각각 배치 사이즈 만큼 가져오기 때문).
전반적인 알고리즘은 다음과 같습니다.
그 이후 위에서 정의한 Loss function으로 학습을 진행합니다.
5. Experiments
우선 CIFAR-10과 SVHN dataset에 대한 labeled data개수에 따른 성능차이입니다. 또한 실험에 사용한 모델은 모두 Wide ResNet-28을 사용하였습니다.
맨 하단의 위치한 supervised learning은 각 사용할 수 있는 data인 50,000개와 73,257개의 데이터로 학습한 결과입니다. 그래프 상으로도 논문에서 제시한 MixMatch 방법론의 성능이 압도적임을 확인할 수 있습니다.
마찬가지로 다양한 dataset과 다양한 방법론에 대한 성능비교 table을 보아도 훨씬 좋은 성능을 보였습니다.
6. Code review
저자가 공개한 코드 중 핵심 알고리즘인 mixup 부분을 위주로 짚고 넘어가겠습니다.
1) dataset 구성
코드에 사용한 데이터는 cifar10으로 train에 250개의 labeled data와 44750개의 unlabeled data를 valid에 5000개를 사용하였습니다.
import numpy as np
from PIL import Image
import torchvision
import torch
class TransformTwice:
def __init__(self, transform):
self.transform = transform
def __call__(self, inp):
out1 = self.transform(inp)
out2 = self.transform(inp)
return out1, out2
def get_cifar10(root, n_labeled,
transform_train=None, transform_val=None,
download=True):
base_dataset = torchvision.datasets.CIFAR10(root, train=True, download=download)
train_labeled_idxs, train_unlabeled_idxs, val_idxs = train_val_split(base_dataset.targets, int(n_labeled/10))
train_labeled_dataset = CIFAR10_labeled(root, train_labeled_idxs, train=True, transform=transform_train)
train_unlabeled_dataset = CIFAR10_unlabeled(root, train_unlabeled_idxs, train=True, transform=TransformTwice(transform_train))
val_dataset = CIFAR10_labeled(root, val_idxs, train=True, transform=transform_val, download=True)
test_dataset = CIFAR10_labeled(root, train=False, transform=transform_val, download=True)
print (f"#Labeled: {len(train_labeled_idxs)} #Unlabeled: {len(train_unlabeled_idxs)} #Val: {len(val_idxs)}")
return train_labeled_dataset, train_unlabeled_dataset, val_dataset, test_dataset
def train_val_split(labels, n_labeled_per_class):
labels = np.array(labels)
train_labeled_idxs = []
train_unlabeled_idxs = []
val_idxs = []
for i in range(10):
idxs = np.where(labels == i)[0]
np.random.shuffle(idxs)
train_labeled_idxs.extend(idxs[:n_labeled_per_class])
train_unlabeled_idxs.extend(idxs[n_labeled_per_class:-500])
val_idxs.extend(idxs[-500:])
np.random.shuffle(train_labeled_idxs)
np.random.shuffle(train_unlabeled_idxs)
np.random.shuffle(val_idxs)
return train_labeled_idxs, train_unlabeled_idxs, val_idxs
cifar10_mean = (0.4914, 0.4822, 0.4465) # equals np.mean(train_set.train_data, axis=(0,1,2))/255
cifar10_std = (0.2471, 0.2435, 0.2616) # equals np.std(train_set.train_data, axis=(0,1,2))/255
def normalise(x, mean=cifar10_mean, std=cifar10_std):
x, mean, std = [np.array(a, np.float32) for a in (x, mean, std)]
x -= mean*255
x *= 1.0/(255*std)
return x
def transpose(x, source='NHWC', target='NCHW'):
return x.transpose([source.index(d) for d in target])
def pad(x, border=4):
return np.pad(x, [(0, 0), (border, border), (border, border)], mode='reflect')
class RandomPadandCrop(object):
"""Crop randomly the image.
Args:
output_size (tuple or int): Desired output size. If int, square crop
is made.
"""
def __init__(self, output_size):
assert isinstance(output_size, (int, tuple))
if isinstance(output_size, int):
self.output_size = (output_size, output_size)
else:
assert len(output_size) == 2
self.output_size = output_size
def __call__(self, x):
x = pad(x, 4)
h, w = x.shape[1:]
new_h, new_w = self.output_size
top = np.random.randint(0, h - new_h)
left = np.random.randint(0, w - new_w)
x = x[:, top: top + new_h, left: left + new_w]
return x
class RandomFlip(object):
"""Flip randomly the image.
"""
def __call__(self, x):
if np.random.rand() < 0.5:
x = x[:, :, ::-1]
return x.copy()
class GaussianNoise(object):
"""Add gaussian noise to the image.
"""
def __call__(self, x):
c, h, w = x.shape
x += np.random.randn(c, h, w) * 0.15
return x
class ToTensor(object):
"""Transform the image to tensor.
"""
def __call__(self, x):
x = torch.from_numpy(x)
return x
class CIFAR10_labeled(torchvision.datasets.CIFAR10):
def __init__(self, root, indexs=None, train=True,
transform=None, target_transform=None,
download=False):
super(CIFAR10_labeled, self).__init__(root, train=train,
transform=transform, target_transform=target_transform,
download=download)
if indexs is not None:
self.data = self.data[indexs]
self.targets = np.array(self.targets)[indexs]
self.data = transpose(normalise(self.data))
def __getitem__(self, index):
"""
Args:
index (int): Index
Returns:
tuple: (image, target) where target is index of the target class.
"""
img, target = self.data[index], self.targets[index]
if self.transform is not None:
img = self.transform(img)
if self.target_transform is not None:
target = self.target_transform(target)
return img, target
class CIFAR10_unlabeled(CIFAR10_labeled):
def __init__(self, root, indexs, train=True,
transform=None, target_transform=None,
download=False):
super(CIFAR10_unlabeled, self).__init__(root, indexs, train=train,
transform=transform, target_transform=target_transform,
download=download)
self.targets = np.array([-1 for i in range(len(self.targets))])
2) model
논문에서 사용한 wide ResNet28 아키텍쳐입니다.
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
class BasicBlock(nn.Module):
def __init__(self, in_planes, out_planes, stride, dropRate=0.0, activate_before_residual=False):
super(BasicBlock, self).__init__()
self.bn1 = nn.BatchNorm2d(in_planes, momentum=0.001)
self.relu1 = nn.LeakyReLU(negative_slope=0.1, inplace=True)
self.conv1 = nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(out_planes, momentum=0.001)
self.relu2 = nn.LeakyReLU(negative_slope=0.1, inplace=True)
self.conv2 = nn.Conv2d(out_planes, out_planes, kernel_size=3, stride=1,
padding=1, bias=False)
self.droprate = dropRate
self.equalInOut = (in_planes == out_planes)
self.convShortcut = (not self.equalInOut) and nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride,
padding=0, bias=False) or None
self.activate_before_residual = activate_before_residual
def forward(self, x):
if not self.equalInOut and self.activate_before_residual == True:
x = self.relu1(self.bn1(x))
else:
out = self.relu1(self.bn1(x))
out = self.relu2(self.bn2(self.conv1(out if self.equalInOut else x)))
if self.droprate > 0:
out = F.dropout(out, p=self.droprate, training=self.training)
out = self.conv2(out)
return torch.add(x if self.equalInOut else self.convShortcut(x), out)
class NetworkBlock(nn.Module):
def __init__(self, nb_layers, in_planes, out_planes, block, stride, dropRate=0.0, activate_before_residual=False):
super(NetworkBlock, self).__init__()
self.layer = self._make_layer(block, in_planes, out_planes, nb_layers, stride, dropRate, activate_before_residual)
def _make_layer(self, block, in_planes, out_planes, nb_layers, stride, dropRate, activate_before_residual):
layers = []
for i in range(int(nb_layers)):
layers.append(block(i == 0 and in_planes or out_planes, out_planes, i == 0 and stride or 1, dropRate, activate_before_residual))
return nn.Sequential(*layers)
def forward(self, x):
return self.layer(x)
class WideResNet(nn.Module):
def __init__(self, num_classes, depth=28, widen_factor=2, dropRate=0.0):
super(WideResNet, self).__init__()
nChannels = [16, 16*widen_factor, 32*widen_factor, 64*widen_factor]
assert((depth - 4) % 6 == 0)
n = (depth - 4) / 6
block = BasicBlock
# 1st conv before any network block
self.conv1 = nn.Conv2d(3, nChannels[0], kernel_size=3, stride=1,
padding=1, bias=False)
# 1st block
self.block1 = NetworkBlock(n, nChannels[0], nChannels[1], block, 1, dropRate, activate_before_residual=True)
# 2nd block
self.block2 = NetworkBlock(n, nChannels[1], nChannels[2], block, 2, dropRate)
# 3rd block
self.block3 = NetworkBlock(n, nChannels[2], nChannels[3], block, 2, dropRate)
# global average pooling and classifier
self.bn1 = nn.BatchNorm2d(nChannels[3], momentum=0.001)
self.relu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
self.fc = nn.Linear(nChannels[3], num_classes)
self.nChannels = nChannels[3]
for m in self.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
elif isinstance(m, nn.Linear):
nn.init.xavier_normal_(m.weight.data)
m.bias.data.zero_()
def forward(self, x):
out = self.conv1(x)
out = self.block1(out)
out = self.block2(out)
out = self.block3(out)
out = self.relu(self.bn1(out))
out = F.avg_pool2d(out, 8)
out = out.view(-1, self.nChannels)
return self.fc(out)
3) loss function
class SemiLoss(object):
def __call__(self, outputs_x, targets_x, outputs_u, targets_u, epoch):
probs_u = torch.softmax(outputs_u, dim=1)
Lx = -torch.mean(torch.sum(F.log_softmax(outputs_x, dim=1) * targets_x, dim=1))
Lu = torch.mean((probs_u - targets_u)**2)
return Lx, Lu, args.lambda_u * linear_rampup(epoch)
4) mixup 알고리즘
위에 정의한 수식처럼 mixup 방법을 이용하여 새롭게 mixed_input과 mixed_target을 만들어 다시 모델에 집어넣어 loss식에 반영하는 구조를 볼 수 있습니다.
def train(labeled_trainloader, unlabeled_trainloader, model, optimizer, ema_optimizer, criterion, epoch, use_cuda):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
losses_x = AverageMeter()
losses_u = AverageMeter()
ws = AverageMeter()
end = time.time()
bar = Bar('Training', max=args.train_iteration)
labeled_train_iter = iter(labeled_trainloader)
unlabeled_train_iter = iter(unlabeled_trainloader)
model.train()
for batch_idx in range(args.train_iteration):
try:
inputs_x, targets_x = labeled_train_iter.next()
except:
labeled_train_iter = iter(labeled_trainloader)
inputs_x, targets_x = labeled_train_iter.next()
try:
(inputs_u, inputs_u2), _ = unlabeled_train_iter.next()
except:
unlabeled_train_iter = iter(unlabeled_trainloader)
(inputs_u, inputs_u2), _ = unlabeled_train_iter.next()
# measure data loading time
data_time.update(time.time() - end)
batch_size = inputs_x.size(0)
# Transform label to one-hot
targets_x = torch.zeros(batch_size, 10).scatter_(1, targets_x.view(-1,1).long(), 1)
if use_cuda:
inputs_x, targets_x = inputs_x.cuda(), targets_x.cuda(non_blocking=True)
inputs_u = inputs_u.cuda()
inputs_u2 = inputs_u2.cuda()
with torch.no_grad():
# compute guessed labels of unlabel samples
outputs_u = model(inputs_u)
outputs_u2 = model(inputs_u2)
p = (torch.softmax(outputs_u, dim=1) + torch.softmax(outputs_u2, dim=1)) / 2
pt = p**(1/args.T)
targets_u = pt / pt.sum(dim=1, keepdim=True)
targets_u = targets_u.detach()
'''
mixup 부분
unlabeled에 sharpen을 적용하여 얻은 targets_u를 concat하여 사용
'''
all_inputs = torch.cat([inputs_x, inputs_u, inputs_u2], dim=0)
all_targets = torch.cat([targets_x, targets_u, targets_u], dim=0)
l = np.random.beta(args.alpha, args.alpha)
l = max(l, 1-l)
idx = torch.randperm(all_inputs.size(0))
input_a, input_b = all_inputs, all_inputs[idx]
target_a, target_b = all_targets, all_targets[idx]
mixed_input = l * input_a + (1 - l) * input_b
mixed_target = l * target_a + (1 - l) * target_b
# interleave labeled and unlabed samples between batches to get correct batchnorm calculation
mixed_input = list(torch.split(mixed_input, batch_size))
mixed_input = interleave(mixed_input, batch_size)
logits = [model(mixed_input[0])]
for input in mixed_input[1:]:
logits.append(model(input))
# put interleaved samples back
logits = interleave(logits, batch_size)
logits_x = logits[0]
logits_u = torch.cat(logits[1:], dim=0)
Lx, Lu, w = criterion(logits_x, mixed_target[:batch_size], logits_u, mixed_target[batch_size:], epoch+batch_idx/args.train_iteration)
loss = Lx + w * Lu
# record loss
losses.update(loss.item(), inputs_x.size(0))
losses_x.update(Lx.item(), inputs_x.size(0))
losses_u.update(Lu.item(), inputs_x.size(0))
ws.update(w, inputs_x.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
ema_optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# plot progress
bar.suffix = '({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | Loss_x: {loss_x:.4f} | Loss_u: {loss_u:.4f} | W: {w:.4f}'.format(
batch=batch_idx + 1,
size=args.train_iteration,
data=data_time.avg,
bt=batch_time.avg,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
loss_x=losses_x.avg,
loss_u=losses_u.avg,
w=ws.avg,
)
bar.next()
bar.finish()
return (losses.avg, losses_x.avg, losses_u.avg,)
5) 학습
from __future__ import print_function
import argparse
import os
import shutil
import time
import random
import numpy as np
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data as data
import torchvision.transforms as transforms
import torch.nn.functional as F
import models.wideresnet as models
import dataset.cifar10 as dataset
from utils import Bar, Logger, AverageMeter, accuracy, mkdir_p, savefig
from tensorboardX import SummaryWriter
parser = argparse.ArgumentParser(description='PyTorch MixMatch Training')
# Optimization options
parser.add_argument('--epochs', default=1024, type=int, metavar='N',
help='number of total epochs to run')
parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
help='manual epoch number (useful on restarts)')
parser.add_argument('--batch-size', default=64, type=int, metavar='N',
help='train batchsize')
parser.add_argument('--lr', '--learning-rate', default=0.002, type=float,
metavar='LR', help='initial learning rate')
# Checkpoints
parser.add_argument('--resume', default='', type=str, metavar='PATH',
help='path to latest checkpoint (default: none)')
# Miscs
parser.add_argument('--manualSeed', type=int, default=0, help='manual seed')
#Device options
parser.add_argument('--gpu', default='0', type=str,
help='id(s) for CUDA_VISIBLE_DEVICES')
#Method options
parser.add_argument('--n-labeled', type=int, default=250,
help='Number of labeled data')
parser.add_argument('--train-iteration', type=int, default=1024,
help='Number of iteration per epoch')
parser.add_argument('--out', default='result',
help='Directory to output the result')
parser.add_argument('--alpha', default=0.75, type=float)
parser.add_argument('--lambda-u', default=75, type=float)
parser.add_argument('--T', default=0.5, type=float)
parser.add_argument('--ema-decay', default=0.999, type=float)
args = parser.parse_args()
state = {k: v for k, v in args._get_kwargs()}
# Use CUDA
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
use_cuda = torch.cuda.is_available()
# Random seed
if args.manualSeed is None:
args.manualSeed = random.randint(1, 10000)
np.random.seed(args.manualSeed)
best_acc = 0 # best test accuracy
def main():
global best_acc
if not os.path.isdir(args.out):
mkdir_p(args.out)
# Data
print(f'==> Preparing cifar10')
transform_train = transforms.Compose([
dataset.RandomPadandCrop(32),
dataset.RandomFlip(),
dataset.ToTensor(),
])
transform_val = transforms.Compose([
dataset.ToTensor(),
])
train_labeled_set, train_unlabeled_set, val_set, test_set = dataset.get_cifar10('./data', args.n_labeled, transform_train=transform_train, transform_val=transform_val)
labeled_trainloader = data.DataLoader(train_labeled_set, batch_size=args.batch_size, shuffle=True, num_workers=0, drop_last=True)
unlabeled_trainloader = data.DataLoader(train_unlabeled_set, batch_size=args.batch_size, shuffle=True, num_workers=0, drop_last=True)
val_loader = data.DataLoader(val_set, batch_size=args.batch_size, shuffle=False, num_workers=0)
test_loader = data.DataLoader(test_set, batch_size=args.batch_size, shuffle=False, num_workers=0)
# Model
print("==> creating WRN-28-2")
def create_model(ema=False):
model = models.WideResNet(num_classes=10)
model = model.cuda()
if ema:
for param in model.parameters():
param.detach_()
return model
model = create_model()
ema_model = create_model(ema=True)
cudnn.benchmark = True
print(' Total params: %.2fM' % (sum(p.numel() for p in model.parameters())/1000000.0))
train_criterion = SemiLoss()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=args.lr)
ema_optimizer= WeightEMA(model, ema_model, alpha=args.ema_decay)
start_epoch = 0
# Resume
title = 'noisy-cifar-10'
if args.resume:
# Load checkpoint.
print('==> Resuming from checkpoint..')
assert os.path.isfile(args.resume), 'Error: no checkpoint directory found!'
args.out = os.path.dirname(args.resume)
checkpoint = torch.load(args.resume)
best_acc = checkpoint['best_acc']
start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
ema_model.load_state_dict(checkpoint['ema_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logger = Logger(os.path.join(args.out, 'log.txt'), title=title, resume=True)
else:
logger = Logger(os.path.join(args.out, 'log.txt'), title=title)
logger.set_names(['Train Loss', 'Train Loss X', 'Train Loss U', 'Valid Loss', 'Valid Acc.', 'Test Loss', 'Test Acc.'])
writer = SummaryWriter(args.out)
step = 0
test_accs = []
# Train and val
for epoch in range(start_epoch, args.epochs):
print('\nEpoch: [%d | %d] LR: %f' % (epoch + 1, args.epochs, state['lr']))
train_loss, train_loss_x, train_loss_u = train(labeled_trainloader, unlabeled_trainloader, model, optimizer, ema_optimizer, train_criterion, epoch, use_cuda)
_, train_acc = validate(labeled_trainloader, ema_model, criterion, epoch, use_cuda, mode='Train Stats')
val_loss, val_acc = validate(val_loader, ema_model, criterion, epoch, use_cuda, mode='Valid Stats')
test_loss, test_acc = validate(test_loader, ema_model, criterion, epoch, use_cuda, mode='Test Stats ')
step = args.train_iteration * (epoch + 1)
writer.add_scalar('losses/train_loss', train_loss, step)
writer.add_scalar('losses/valid_loss', val_loss, step)
writer.add_scalar('losses/test_loss', test_loss, step)
writer.add_scalar('accuracy/train_acc', train_acc, step)
writer.add_scalar('accuracy/val_acc', val_acc, step)
writer.add_scalar('accuracy/test_acc', test_acc, step)
# append logger file
logger.append([train_loss, train_loss_x, train_loss_u, val_loss, val_acc, test_loss, test_acc])
# save model
is_best = val_acc > best_acc
best_acc = max(val_acc, best_acc)
save_checkpoint({
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'ema_state_dict': ema_model.state_dict(),
'acc': val_acc,
'best_acc': best_acc,
'optimizer' : optimizer.state_dict(),
}, is_best)
test_accs.append(test_acc)
logger.close()
writer.close()
print('Best acc:')
print(best_acc)
print('Mean acc:')
print(np.mean(test_accs[-20:]))
def validate(valloader, model, criterion, epoch, use_cuda, mode):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
# switch to evaluate mode
model.eval()
end = time.time()
bar = Bar(f'{mode}', max=len(valloader))
with torch.no_grad():
for batch_idx, (inputs, targets) in enumerate(valloader):
# measure data loading time
data_time.update(time.time() - end)
if use_cuda:
inputs, targets = inputs.cuda(), targets.cuda(non_blocking=True)
# compute output
outputs = model(inputs)
loss = criterion(outputs, targets)
# measure accuracy and record loss
prec1, prec5 = accuracy(outputs, targets, topk=(1, 5))
losses.update(loss.item(), inputs.size(0))
top1.update(prec1.item(), inputs.size(0))
top5.update(prec5.item(), inputs.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# plot progress
bar.suffix = '({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | top1: {top1: .4f} | top5: {top5: .4f}'.format(
batch=batch_idx + 1,
size=len(valloader),
data=data_time.avg,
bt=batch_time.avg,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
top1=top1.avg,
top5=top5.avg,
)
bar.next()
bar.finish()
return (losses.avg, top1.avg)
def save_checkpoint(state, is_best, checkpoint=args.out, filename='checkpoint.pth.tar'):
filepath = os.path.join(checkpoint, filename)
torch.save(state, filepath)
if is_best:
shutil.copyfile(filepath, os.path.join(checkpoint, 'model_best.pth.tar'))
def linear_rampup(current, rampup_length=args.epochs):
if rampup_length == 0:
return 1.0
else:
current = np.clip(current / rampup_length, 0.0, 1.0)
return float(current)
class WeightEMA(object):
def __init__(self, model, ema_model, alpha=0.999):
self.model = model
self.ema_model = ema_model
self.alpha = alpha
self.params = list(model.state_dict().values())
self.ema_params = list(ema_model.state_dict().values())
self.wd = 0.02 * args.lr
for param, ema_param in zip(self.params, self.ema_params):
param.data.copy_(ema_param.data)
def step(self):
one_minus_alpha = 1.0 - self.alpha
for param, ema_param in zip(self.params, self.ema_params):
if ema_param.dtype==torch.float32:
ema_param.mul_(self.alpha)
ema_param.add_(param * one_minus_alpha)
# customized weight decay
param.mul_(1 - self.wd)
def interleave_offsets(batch, nu):
groups = [batch // (nu + 1)] * (nu + 1)
for x in range(batch - sum(groups)):
groups[-x - 1] += 1
offsets = [0]
for g in groups:
offsets.append(offsets[-1] + g)
assert offsets[-1] == batch
return offsets
def interleave(xy, batch):
nu = len(xy) - 1
offsets = interleave_offsets(batch, nu)
xy = [[v[offsets[p]:offsets[p + 1]] for p in range(nu + 1)] for v in xy]
for i in range(1, nu + 1):
xy[0][i], xy[i][i] = xy[i][i], xy[0][i]
return [torch.cat(v, dim=0) for v in xy]
if __name__ == '__main__':
main()
6) 결과
250개의 label을 주었을 때의 결과만 비교해 보았습니다. 기존 논문에서 주장하는 성능과 큰 차이가 없었습니다.
