안녕하세요.

 

이번 글에서는 저번 글에서 리뷰했던 Wide Residual Networks(WRN)의 code review를 진행해보겠습니다.

 

ResNet과 유사한 형태이기 때문에, ResNet를 이해하셨다면 조금 더 쉽게 보실 수 있을 것 같습니다.

 

해당 논문에 대한 paper review는 다음 주소에서 확인하실 수 있습니다.

 

cumulu-s.tistory.com/35

 

7. Wide Residual Networks(WRN) - paper review

안녕하세요. 오늘은 ResNet을 발전시켜서 만들어진 Wide Residual Networks(WRN)에 대한 내용을 정리해보려고 합니다. 논문 주소: arxiv.org/abs/1605.07146 Wide Residual Networks Deep residual networks were..

cumulu-s.tistory.com

 

시작해보겠습니다.

 

 

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from tqdm.auto import tqdm
import torchvision
import torchvision.transforms as transforms
import datetime
import os
from torch.utils.tensorboard import SummaryWriter
import shutil

먼저, 필요한 library를 가져옵니다.

 

 

device = 'cuda' if torch.cuda.is_available() else 'cpu'

# Data
print('==> Preparing data..')
transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])

trainset = torchvision.datasets.CIFAR10(
    root='/content/drive/MyDrive/CIFAR10', train=True, download=True, transform=transform_train)
trainloader = torch.utils.data.DataLoader(
    trainset, batch_size=128, shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(
    root='/content/drive/MyDrive/CIFAR10', train=False, download=True, transform=transform_test)
testloader = torch.utils.data.DataLoader(
    testset, batch_size=100, shuffle=False, num_workers=2)

current_time = datetime.datetime.now() + datetime.timedelta(hours= 9)
current_time = current_time.strftime('%Y-%m-%d-%H:%M')

saved_loc = os.path.join('/content/drive/MyDrive/WRN_Result', current_time)
if os.path.exists(saved_loc):
    shutil.rmtree(saved_loc)
os.mkdir(saved_loc)

print("결과 저장 위치: ", saved_loc)

writer = SummaryWriter(saved_loc)

데이터를 들고 올 준비를 해줍니다.

 

 

class BasicBlock(nn.Module):
    def __init__(self, in_planes, out_planes, stride, dropRate=0.0):
        super(BasicBlock, self).__init__()
        self.bn1 = nn.BatchNorm2d(in_planes)
        self.relu1 = nn.ReLU(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)
        self.relu2 = nn.ReLU(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
    def forward(self, x):
        if not self.equalInOut:
            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):
        super(NetworkBlock, self).__init__()
        self.layer = self._make_layer(block, in_planes, out_planes, nb_layers, stride, dropRate)
    def _make_layer(self, block, in_planes, out_planes, nb_layers, stride, dropRate):
        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))
        return nn.Sequential(*layers)
    def forward(self, x):
        return self.layer(x)

class WideResNet(nn.Module):
    def __init__(self, depth, num_classes, widen_factor=1, dropRate=0.0):
        super(WideResNet, self).__init__()
        nChannels = [16, 16*widen_factor, 32*widen_factor, 64*widen_factor]
        assert((depth - 4) % 6 == 0) # depth must be 6n + 4
        n = (depth - 4) / 6

        # 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], BasicBlock, 1, dropRate)
        # 2nd block
        self.block2 = NetworkBlock(n, nChannels[1], nChannels[2], BasicBlock, 2, dropRate)
        # 3rd block
        self.block3 = NetworkBlock(n, nChannels[2], nChannels[3], BasicBlock, 2, dropRate)
        # global average pooling and classifier
        self.bn1 = nn.BatchNorm2d(nChannels[3])
        self.relu = nn.ReLU(inplace=True)
        self.fc = nn.Linear(nChannels[3], num_classes)
        self.nChannels = nChannels[3]


        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):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                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, 1, 0)
        out = out.view(-1, self.nChannels)
        return self.fc(out)

 

WideResNet이라는 이름으로 WRN를 나타내는 class를 만들어줍니다. 

 

input으로는 depth, class 개수, widen_factor, dropout rate를 받게 됩니다.

 

Channel의 개수는 16, 16*k, 32*k, 64*k로 구성되며, k는 이전 논문 review에서도 언급되었듯이 widening factor를 나타냅니다.

 

이를 통해 기존 original residual network보다 더 넓은 형태의 residual network를 구성할 수 있습니다.

 

depth의 경우는 반드시 6n + 4의 값이 되어야 하는데, 이는 다음 표를 보고 이해해볼 수 있습니다.

 

먼저, conv1에서 1개, 그리고 conv2, conv3, conv4에서 shortcut connection 형태로 layer가 들어가게 되어 기본적으로 4개의 layer가 있어야 합니다.

 

그리고 6n이 되는 이유는, conv2에서 3x3 짜리 2개, conv3에서 3x3짜리 2개, conv4에서 3x3짜리 2개로 총 6개가 기본적인 setting이므로, 이 setting보다 몇 배 더 많은 layer를 구성할 것인지를 결정하기 위해 6n가 되게 됩니다. 

 

만약 layer의 총개수가 16이라면, 6*2 + 4가 되며, 이는 conv2 block 내에서의 layer가 4개, conv3 block 내에서의 layer가 4개, conv4 block 내에서의 layer가 4개가 됩니다. 

 

즉, n이라는 것은 conv block 내에서 3x3 conv layer의 개수에 곱해지는 계수라고 생각할 수 있습니다.

 

 

 

위의 Table 1의 structure가 코드로 구현되어 있음을 확인할 수 있는데요.

 

WRN의 전체적인 흐름부터 살펴보겠습니다.

 

    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, 1, 0)
        out = out.view(-1, self.nChannels)
        return self.fc(out)

(WideResNet에서 이쪽 부분을 확인해주시면 됩니다!)

 

 

먼저 self.conv1은 3x3 conv이며, 16 channel이 output이 되도록 연산이 되게 됩니다.

 

self.block1, self.block2, self.block3는 위 Table 1에서 conv2, conv3, conv4에 대응됩니다.

 

self.block에 대해서는 조금 있다가 자세히 살펴보도록 할 예정입니다.

 

self.block3까지 지난 결과를 BN과 ReLU를 통과시키고, average pooling을 적용해줍니다. 

 

그리고 이를 [batch_size, 64*widen_factor]의 shape로 만들어주고, 이를 마지막 fully connected layer에 연결하여

 

dataset의 class 각각에 대한 probability를 정답으로 내게 됩니다.

 

 

 

이전에 paper review를 진행할 때, WRN의 contribution 중 하나가 바로 batch normalization, activation, convolutional layer의 순서에 대한 연구를 진행한 것이었고, 따라서 어떤 순서로 연산할지 변경되었다는 얘기를 했었습니다.

 

그렇다면, 여기서 ResNet과 Wide Residual Network는 어떻게 다른지 한번 살펴보겠습니다.

 

class ResNet(nn.Module):
    def __init__(self, block, num_blocks, num_classes=10):
        super(ResNet, self).__init__()
        self.in_planes = 64

        self.conv1 = nn.Conv2d(3, 64, kernel_size=3,
                               stride=1, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
        self.linear = nn.Linear(512*block.expansion, num_classes)

    def _make_layer(self, block, planes, num_blocks, stride):
        strides = [stride] + [1]*(num_blocks-1)
        layers = []
        for stride in strides:
            layers.append(block(self.in_planes, planes, stride))
            self.in_planes = planes * block.expansion
        return nn.Sequential(*layers)

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        out = F.avg_pool2d(out, 4)
        out = out.view(out.size(0), -1)
        out = self.linear(out)
        return out

 

이전 ResNet code review를 할 때 사용했던 코드입니다.

 

ResNet에서는 먼저 input 데이터를 받아 conv1 -> BN -> relu를 거치고 conv block을 거친 후, average pooling을 하고, 이를 fully connected layer에 연결하는 순서로 이루어집니다.

 

하지만, Wide residual network는 input 데이터를 받아 conv1을 거치고, conv block을 거친 후, BN -> relu를 거치고 average pooling을 한 후 이를 fully connected layer에 연결하는 순서로 이루어집니다.

 

전체적인 흐름상 batch normalization과 ReLU의 적용 위치가 바뀌었네요.

 

그리고 conv block이 ResNet에서는 4개였으나, Wide residual network에서는 3개라는 점도 하나의 차이가 될 수 있겠네요.

 

conv block 1개당 stride = 2가 적용되므로, 따라서 ResNet은 conv block이 모두 진행되었을 때 resolution이 4x4지만, wide residual network는 resolution이 8x8이 됩니다.

 

 

 

다음으로는 conv block에 대해서 세부적으로 살펴보겠습니다.

 

self.block1 = NetworkBlock(n, nChannels[0], nChannels[1], BasicBlock, 1, dropRate)

 

먼저 첫 번째 conv block을 살펴봅니다.

 

간단한 예시를 들기 위해서, n = 2 / nChannels[0] = 16 / nChannels[1] = 64 / dropRate = 0.3으로 지정한다고 가정하겠습니다.

 

class NetworkBlock(nn.Module):
    def __init__(self, nb_layers, in_planes, out_planes, block, stride, dropRate=0.0):
        super(NetworkBlock, self).__init__()
        self.layer = self._make_layer(block, in_planes, out_planes, nb_layers, stride, dropRate)
    def _make_layer(self, block, in_planes, out_planes, nb_layers, stride, dropRate):
        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))
        return nn.Sequential(*layers)
    def forward(self, x):
        return self.layer(x)

 

따라서, 우리가 만들어낸 것은 NetworkBlock(nb_layers = 2, in_planes = 16, out_planes = 64, block = BasicBlock, stride = 1, dropRate= 0.3)이 됩니다.

 

이를 통해서 만들어지는 것은 다음과 같습니다.

 

def _make_layer(self, block = BasicBlock, in_planes=16, out_planes=64, nb_layers=2, stride=1, dropRate=0.3):
        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))
        return nn.Sequential(*layers)

 

layer의 개수만큼 for문을 돌게 되고, 이를 통해서 layers = []에 block이 append 되는 구조입니다.

 

nb_layers = 2이므로, block이 2개 쌓이게 되겠죠?

 

여기서 조금 생소한 코드가 등장하는데요, A and B or C 형태의 문법이 사용되었습니다.

 

이는 B if A else C와 똑같은 의미입니다.

 

즉, BasicBlock(in_planes if i == 0 else out_planes, out_planes, stride if i == 0 else 1, dropRate)가 만들어집니다.

 

layers = [BasicBlock(16, 64, 1, 0.3), BasicBlock(64, 64, 1, 0.3)]가 될 것입니다.

 

다음으로는 BasicBlock에 대해서 살펴봐야겠죠?

 

class BasicBlock(nn.Module):
    def __init__(self, in_planes, out_planes, stride, dropRate=0.0):
        super(BasicBlock, self).__init__()
        self.bn1 = nn.BatchNorm2d(in_planes)
        self.relu1 = nn.ReLU(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)
        self.relu2 = nn.ReLU(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
    def forward(self, x):
        if not self.equalInOut:
            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)

 

BasicBlock(16, 64, 1, 0.3)이라고 가정하고 생각해보겠습니다.

 

여기서 먼저 따져봐야 할 점은 input의 shape와 output의 shape가 같은지입니다. 이를 확인한 후, self.equalInOut이라는 변수로 boolean 값으로 저장합니다.

 

BasicBlock(16, 64, 1, 0.3)이라면 input의 shape와 output의 shape가 다른 경우이므로, self.equalInOut은 False가 됩니다.

 

self.convShortcut의 경우는 self.equalInOut이 False이면 nn.conv2d를 적용하고, True이면 None이 됩니다.

 

다음으로는 forward 부분을 보면 됩니다.

 

if not self.equalInOut이므로, 이 if문에 들어가게 됩니다.

 

따라서 x = self.relu1(self.bn1(x))이 됩니다.

 

그러고 나서 out = self.relu2(self.bn2(self.conv1(x)))이 되고

 

droprate = 0.3이므로, 여기에 dropout을 적용하여 out = F.dropout(out, p = self.droprate, training = self.training)을 적용합니다.

 

training = self.training을 적용하는 이유는, 현재 모델이 train모드인지 eval 모드인지에 따라서 dropout을 사용할지 아닐지를 결정하게 됩니다.

 

왜냐하면 학습이 된 상태에서 추론을 하는 경우에는 dropout을 사용하지 않아야 하기 때문이죠.

 

그리고 나서 out = self.conv2(out)을 해주고

 

torch.add(self.convShortcut(x))을 해준 것을 return 해줍니다.

 

이를 그림으로 한번 그려보자면, 다음과 같습니다.

 

 

k는 kernel size이며 s는 stride, p는 padding을 나타냅니다.

 

본 그림을 통해서 BasicBlock의 구동 방식을 이해해볼 수 있을 것이라고 생각이 들고요.

 

만약 stride = 2인 경우는 self.conv1에서 s = 2로 변하고, 그 이후로 가로와 세로의 길이가 /2로 줄어듭니다.

 

그리고 input과 output의 channel 수가 동일한 경우는 위의 그림에서 self.conv1 = nn.Conv2d(16, 64, ~)을 통과하면서 channel수가 변하는 것이 아닌, 그대로 channel 수가 유지되게 됩니다.

 

self.convShortcut 또한 필요가 없어지기 때문에, 이는 None이 되고 위의 def forward에서 return 쪽을 확인해보시게 되면, input인 x를 그대로 torch.add에 넣어주는 것을 확인할 수 있습니다. 

 

여기까지 WRN의 모델에 대한 설명은 마무리가 되었습니다.

 

 

net = WideResNet(40, 10, 2, 0.3) # WRN-40-2

net = net.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.1,
                      momentum=0.9, weight_decay=0.0005)
scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones = [60, 120, 160], gamma = 0.2)

best_acc = 0

 

WideResNet의 첫 번째 인자는 layer의 수, 두 번째 인자는 class의 수, 세 번째 인자는 widening factor, 네 번째 인자는 dropout rate입니다.

 

optimizer는 논문에 나온 그대로 lr = 0.1, momentum = 0.9, weight_decay = 0.0005를 적용하였으며

 

scheduler는 MultiStepLR을 이용하여 구현하였습니다. milestones을 통해서 어떤 지점에서 scheduler를 적용할지 입력하고, gamma는 얼마큼 감소시킬지 값을 알려줍니다.

 

처음 learning rate는 0.1로, 그리고 60, 120, 160 에폭 일 때 각각 0.2씩 감소하라는 내용을 그대로 반영하였습니다.

 

 

def train(epoch):
    print('\nEpoch: %d' % epoch)
    net.train()
    train_loss = 0
    correct = 0
    total = 0
    batch_count = 0
    for batch_idx, (inputs, targets) in enumerate(trainloader):
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()
        batch_count += 1

    print("Train Loss : {:.3f} | Train Acc: {:.3f}".format(train_loss / batch_count, 100.*correct/total))
    final_loss = train_loss / batch_count
    final_acc = 100.*correct / total
    return final_loss, final_acc

 

inputs로 이미지 데이터를, targets로 label을 받아오고

 

이미지를 신경망에 투입하여 outputs을 만들고

 

outputs과 targets 과의 cross entropy를 계산해서 loss로 만듭니다.

 

이 값은 backpropagation 시키고, train_loss에 더해줍니다.

 

output.max(1)을 적용하면, row 기준으로 output에서 가장 높은 값을 가지는 index와 예측값을 뽑아줍니다.

 

예측값은 필요 없으므로 _로 처리하고, 예측된 index를 predicted로 저장합니다.

 

그리고 predicted.eq(targets)를 통해서 label과 예측된 label이 몇 개나 같은지를 boolean 값으로 계산합니다.

 

그리고 sum()을 해주면, 배치 사이즈 기준으로 몇 개나 label와 일치하는지를 값으로 계산할 수 있습니다. 

 

 

def test(epoch):
    global best_acc
    net.eval()
    test_loss = 0
    correct = 0
    total = 0
    batch_count = 0
    with torch.no_grad():
        for batch_idx, (inputs, targets) in enumerate(testloader):
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = net(inputs)
            loss = criterion(outputs, targets)

            test_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
            batch_count += 1


    print("Test Loss : {:.3f} | Test Acc: {:.3f}".format(test_loss / batch_count, 100.*correct/total))

    # Save checkpoint.
    acc = 100.*correct/total
    if acc > best_acc:
        print('Saving..')
        state = {
            'net': net.state_dict(),
            'acc': acc,
            'epoch': epoch,
        }
        if not os.path.isdir(os.path.join(saved_loc, 'checkpoint')):
            os.mkdir(os.path.join(saved_loc, 'checkpoint'))
        torch.save(state, os.path.join(saved_loc, 'checkpoint/ckpt.pth'))
        best_acc = acc
    
    final_loss = test_loss / batch_count
    return final_loss, acc

 

위의 내용은 동일하니 설명을 생략하고,

 

accuracy의 경우 best_acc를 계속 갱신하면서, best_acc보다 더 높은 accuracy가 나오면 이를 best_acc로 저장하면서 동시에 모델과 accuracy 값, 그 당시 epoch을 저장합니다.

 

이를 통해서 test 기준으로 가장 높게 나온 네트워크의 가중치와 accuracy 값을 저장할 수 있게 됩니다.

 

 

for epoch in tqdm(range(200)):
    train_loss, train_acc = train(epoch)
    test_loss, test_acc = test(epoch)

    writer.add_scalars('Loss', {"Train Loss" : train_loss, "Test Loss" : test_loss}, epoch)
    writer.add_scalars('Accuracy', {"Train acc" : train_acc, "test acc" : test_acc}, epoch)
    
    scheduler.step()

writer.close()

학습은 200 epochs 동안 진행하고, writer.add_scalars는 tensorboard에 값을 저장하는 코드입니다.

 

이를 통해서 train_loss, test_loss, train_acc, test_acc 그래프를 그릴 수 있게 됩니다.

 

그리고 앞에서 scheduler를 지정했기 때문에, 1 epoch이 지날 때마다 scheduler.step()을 적용하여 원하는 지점에서 learning rate가 변하도록 지정해줍니다. 

 

 

 

이를 통해 얻게 되는 Accuracy 그래프는 다음과 같습니다.

 

 

 

중간에 learning rate가 변했다 보니, 약간 각진 그래프가 나오게 되네요.

 

 

 

loss graph도 마찬가지로 잘 수렴하는 모습을 보여주고 있습니다.

 

 

 

여기까지 WRN의 code review를 해보았습니다.

 

해당 글에서 다룬 모든 코드는 제 Github에서 확인하실 수 있습니다.

 

github.com/PeterKim1/paper_code_review

 

PeterKim1/paper_code_review

paper review with codes. Contribute to PeterKim1/paper_code_review development by creating an account on GitHub.

github.com

 

다음에는 또 새로운 논문으로 찾아오도록 하겠습니다.

 

감사합니다.

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

+ Recent posts