CHAPTER5 이미지 처리 능력이 탁월한 CNN 2

EPOCHS = 300
BATCH_SIZE = 128

# 데이터 불러오기
train_loader = torch.utils.data.DataLoader(
    datasets.CIFAR10('./.data',
                    train = True,
                    download = True,
                    transform = transforms.Compose([
                        transforms.RandomCrop(32, padding = 4),
                        transforms.RandomHorizontalFlip(),
                        transforms.ToTensor(),
                        transforms.Normalize((0.5, 0.5, 0.5),
                                            (0.5, 0.5, 0.5))
                    ])),
    batch_size = BATCH_SIZE, shuffle = True
)

test_loader = torch.utils.data.DataLoader(
    datasets.CIFAR10('./.data',
                    train = False,
                    transform = transforms.Compose([
                        transforms.ToTensor(),
                        transforms.Normalize((0.5, 0.5, 0.5),
                                            (0.5, 0.5, 0.5))
                    ])),
    batch_size = BATCH_SIZE, shuffle = True
)

class BasicBlock(nn.Module):
    def __init__(self, in_planes, planes, stride = 1):
        super(BasicBlock, self).__init__()

        # Conv2d(입력 채널 수 = in_planes, 출력 채널 수=planes)
        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size = 3,
                           stride = stride, padding = 1, bias = False)
        self.bn1 = nn.BatchNorm2d(planes) # 배치 정규화
        self.conv2 = nn.Conv2d(planes, planes, kernel_size = 3,
                              stride = 1, padding = 1, bias = False)
        self.bn2 = nn.BatchNorm2d(planes)

        # 입력 채널인 in_planes 를 받아,
        # self.bn2 계층의 출력 크기와 같은 planes 와 더해주는 shortcut 모듈 정의.
        # Sequential은 여러 모듈을 하나의 모듈로 묶는 역할.
        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != planes:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_planes, planes,
                         kernel_size = 1, stride = stride, bias = False),
                nn.BatchNorm2d(planes)
            )

    # 데이터 흐름
    # 입력 x -> 컨볼루션 -> 배치 정규화 -> 활성화 함수
    # -> 입력 x를 shortcut 거쳐 planes 로 크기 같게 한 거 더하기
    # -> Relu 활성화 함수 통과 -> 최종 출력
    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out

class ResNet(nn.Module):
    def __init__(self, num_classes = 10):
        super(ResNet, self).__init__()
        # self._make_layer 함수에서 각 층을 만들 때 전층 채널 출력값 기록용
        self.in_planes = 16 # layer1 입력 채널 개수 = 16

        self.conv1 = nn.Conv2d(3, 16, kernel_size = 3,
                              stride = 1, padding = 1, bias = False)
        self.bn1 = nn.BatchNorm2d(16)
        self.layer1 = self._make_layer(16, 2, stride = 1) # 반환값 (16, 32, 32)
        self.layer2 = self._make_layer(32, 2, stride = 2) # (32, 16, 16)
        self.layer3 = self._make_layer(64, 2, stride = 2) # (64, 8, 8)

        # 평균 풀링으로 텐서에 있는 원소 개수 64개 -> 10개(레이블 개수)
        self.linear = nn.Linear(64, num_classes)

    # 여러 BasicBlock을 모듈 하나로 묶어주는 역할.
    # 반환값은 컨볼루션 계층과 마찬가지로 모듈 취급하면 된다.
    def _make_layer(self, planes, num_blocks, stride):
        strides = [stride] + [1] * (num_blocks - 1)
        layers = []
        for stride in strides:
            # BasicBlock 생성함.
            # layer1 에서는 16 -> 16 인 BasicBlcok 2개
            # layer2 에서는 16 -> 32 인 BasicBlock 1개, 32 -> 32 인 블록 1개
            # layer3 에서는 32 -> 32 1개, 32 -> 64 1개

            # 여기서 16 -> 32 와 32 -> 64 가 이저너 입력을 중간층에 더해 이미지 맥락 보존 역할.
            layers.append(BasicBlock(self.in_planes, planes, stride))
            self.in_planes = planes
        return nn.Sequential(*layers)


    # 데이터 흐름
    # 1. 입력 들어오면 (컨볼루션, 배치정규화, 활성화함수) 통과
    # 2. 정의해둔 BasicBlock 갖고 있는 layer1, 2, 3 통과
    # 3. 나온 값에 평균 풀링 -> 마지막 계층 분류 결과 출력
    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 = F.avg_pool2d(out, 8)
        out = out.view(out.size(0), -1)
        out = self.linear(out)
        return out

# 학습률 감소 기법 사용
# 학습 진행하면서 최적화 함수의 학습률을 점점 낮춰서 더 정교화게 최적화 함.
model = ResNet().to(DEVICE)
optimizer = optim.SGD(model.parameters(), lr = 0.1,
                     momentum = 0.9, weight_decay = 0.0005)

# 이폭마다 호출되며 step_size가 50이니, 50번 호출될 때 학습률에 gamma(=0.1)값을 곱한다.
# 0.1로 시작한 학습률은 50 이폭 후에 0.01이 됨.
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size = 50, gamma = 0.1)

print(model)

# 학습 및 테스트
# 앞과 다른 점 : scheduler.step() 함수로 학습률 낮춰주는 단계 추가

def train(model, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(DEVICE), target.to(DEVICE)
        optimizer.zero_grad()
        output = model(data)
        loss = F.cross_entropy(output, target)
        loss.backward()
        optimizer.step()

def evaluate(model, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(DEVICE), target.to(DEVICE)
            output = model(data)

            test_loss += F.cross_entropy(output, target, reduction = "sum").item()
            pred = output.max(1, keepdim = True)[1]
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    test_accuracy = 100. * correct / len(test_loader.dataset)
    return test_loss, test_accuracy


for epoch in range(1, EPOCHS + 1):
    scheduler.step()
    train(model, train_loader, optimizer, epoch)
    test_loss, test_accuracy = test(model, test_loader)

    print(f"[{epoch}] Test Loss: {test_loss:.4f}, Accuracy: {test_accuracy:.2f}%")

for epoch in range(1, EPOCHS + 1):
    train(model, train_loader, optimizer, epoch)
    scheduler.step()
    ...