파이토치를 활용한 Convolutional AutoEncoder - CIFAR-10

나중에 변경할 경우 코드 내에서 일일이 수정하지 않아도 되게 하기 위해, 배치 사이즈, 학습률과 에폭 개수는 미리 딕셔너리에 설정해 주었음 현재 gpu를 사용할 수 없는 상태라ㅠㅠ 에폭 개수를 대폭 줄여 모델을 학습 시켰는데 cuda사용이 가능했다면 에폭 개수를 100 정도로 했었을 것 같다!

args = {
	'BATCH_SIZE': 5,
        'LEARNING_RATE': 0.001,
        'NUM_EPOCH': 20
        }

 

 

파이토치 공식 홈페이지의 튜토리얼 페이지를 참고하여 데이터 다운로드
https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html#sphx-glr-beginner-blitz-cifar10-tutorial-py

train_set = torchvision.datasets.CIFAR10('./data', train=True, 
					download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(trainset, batch_size=args['BATCH_SIZE'], 
					shuffle=True, num_workers=2)

test_set = torchvision.datasets.CIFAR10('./data', train=False, 
					download=True, transform=transform)
test_loader = torch.utils.data.DataLoader(testset, batch_size=args['BATCH_SIZE'], 
					shuffle=False, num_workers=2)
                    
                    
classes = ['plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']

torchvision.datasets.CIFAR10 안 root='./data'대신 './data'를 사용해 현재 디렉토리에 'data' 이름의 폴더가 생성되고 그 폴더 안에 데이터가 다운로드 되었음

 

 

모델을 구성하기 전에 먼저 5개의 이미지를 살펴보면,

images, labels = next(iter(train_loader))

fig = plt.figure(figsize=(15,3))
for i, img, label in zip(range(args['BATCH_SIZE']), images, labels):
    numpy_img = np.transpose(img.numpy(), (1, 2, 0))
    ax = fig.add_subplot(1, args['BATCH_SIZE'], i+1)
    ax.imshow(numpy_img)
    title = classes[int(label)]
    ax.set_title(title, fontsize=20)
plt.show()

 

Auto-Encoder

class ConvAutoEncoder(nn.Module):
    def __init__(self):
        super(ConvAutoEncoder, self).__init__()
        
        # Encoder
        self.cnn_layer1 = nn.Sequential(
                        nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1),
                        nn.ReLU(),
                         nn.MaxPool2d(2,2))

        self.cnn_layer2 = nn.Sequential(
                                nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1),
                                nn.ReLU(),
                                 nn.MaxPool2d(2,2))

        # Decoder
        self.tran_cnn_layer1 = nn.Sequential(
                        nn.ConvTranspose2d(32, 16, kernel_size = 2, stride = 2, padding=0),
                        nn.ReLU())

        self.tran_cnn_layer2 = nn.Sequential(
                        nn.ConvTranspose2d(16, 3, kernel_size = 2, stride = 2, padding=0),
                        nn.Sigmoid())
            
            
    def forward(self, x):
        output = self.cnn_layer1(x)
        output = self.cnn_layer2(output)
        output = self.tran_cnn_layer1(output)
        output = self.tran_cnn_layer2(output)

        return output

* Validation은 따로 설정하지 않았음

 

model = ConvAutoEncoder()

criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=args['LEARNING_RATE'])

 

Model Train

steps = 0
total_steps = len(train_loader)
for epoch in range(args['NUM_EPOCH']):
    running_loss = 0
    for i, (X_train, _) in enumerate(train_loader):
        steps += 1

        outputs = model(X_train)
        loss = criterion(outputs, X_train)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        running_loss += loss.item()*X_train.shape[0]
        
        if steps % total_steps == 0:
            model.eval()
            print('Epoch: {}/{}'.format(epoch+1, args['NUM_EPOCH']),
                 'Training loss: {:.5f}..'.format(running_loss/total_steps))
            
            steps = 0
            running_loss = 0
            model.train()

 

테스트 세트의 이미지 중 5개를 플랏 해 보면,

test_images, test_labels = next(iter(test_loader))

fig = plt.figure(figsize=(15,3))
for i, img, label in zip(range(args['BATCH_SIZE']), test_images, test_labels):
    numpy_img = np.transpose(img.detach().numpy(), (1, 2, 0))
    ax = fig.add_subplot(1, args['BATCH_SIZE'], i+1)
    ax.imshow(numpy_img)
    title = classes[int(label)]
    ax.set_title(title, fontsize=20)
plt.show()

 

 

이 5개의 이미지가 학습된 모델을 통과한 뒤의 아웃풋

test_output = model(test_images)

fig = plt.figure(figsize=(15,3))
for i, img, label in zip(range(args['BATCH_SIZE']), test_output, test_labels):
    numpy_img = np.transpose(img.detach().numpy(), (1, 2, 0))
    ax = fig.add_subplot(1, args['BATCH_SIZE'], i+1)
    ax.imshow(numpy_img)
    title = classes[int(label)]
    ax.set_title(title, fontsize=20)
plt.show()