<jupyter_start><jupyter_text>PixelCNN ![title](pics/pixelcnn.png)<jupyter_code>import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import os
os.environ["CUDA_VISIBLE_DEVICES"]="0"
# DEVICE의 출력으로 cuda가 나오는지 확인합니다.
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print('DEVICE :', DEVICE)
# convert data to a normalized torch.FloatTensor
transform = transforms.Compose([
    transforms.ToTensor(),
    ])

# MNIST dataset을 준비합니다. 필요시 다운로드 받습니다.
train_data = datasets.MNIST('mnist_data', train=True, download=True, transform=transform)
print(train_data)

# MNIST dataset을 트레이닝과 테스트에 사용할 수 있도록 Data Loader를 만듭니다.
train_loader = torch.utils.data.DataLoader(train_data, batch_size=32)
print('done')
class CausalConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size):
        super().__init__()
        self.kernel_size = kernel_size
        self.weight = nn.Parameter(torch.randn(out_channels, in_channels, kernel_size, kernel_size))
        self.mask = nn.Parameter(torch.ones(1, 1, kernel_size, kernel_size), requires_grad=False)
        self.mask[:, :, kernel_size//2+1:] = 0.
        self.mask[:, :, kernel_size//2, kernel_size//2:] = 0.
        self.bias = nn.Parameter(torch.zeros(out_channels))
        
    def forward(self, x):
        # x : (batch, channels, height, width)
        weight = self.weight * self.mask
        y = F.conv2d(x, weight, self.bias, padding=self.kernel_size//2)
        return y
        
class CausalConv2dLayer(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size):
        super().__init__()
        self.conv = nn.Sequential(CausalConv2d(in_channels, out_channels, kernel_size),
                                  nn.BatchNorm2d(out_channels),
                                  nn.ReLU())
    
    def forward(self, x):
        # x : (batch, channels, height, width)
        y = x + self.conv(x)
        return y
class PixelCNN(nn.Module):
    def __init__(self, in_channels=1, model_channels=128, out_channels=1, kernel_size=5):
        super().__init__()
        self.in_layer = nn.Conv2d(1, model_channels, kernel_size=1)
        self.convs = nn.Sequential(CausalConv2dLayer(model_channels, model_channels, kernel_size),
                                   CausalConv2dLayer(model_channels, model_channels, kernel_size),
                                   CausalConv2dLayer(model_channels, model_channels, kernel_size),
                                   CausalConv2dLayer(model_channels, model_channels, kernel_size),
                                   CausalConv2dLayer(model_channels, model_channels, kernel_size),
                                   CausalConv2dLayer(model_channels, model_channels, kernel_size))
                                   
        self.out_layer = nn.Conv2d(model_channels, out_channels, kernel_size=1)
        
    def forward(self, x):
        # x : (batch, 1, 28, 28)
        y = self.in_layer(x)
        y = self.convs(y)
        y = self.out_layer(y)
        return y
    
    def inference(self, x):
        # x : (batch, 1, 28, 28)
        for i in range(x.size(2)):
            for j in range(x.size(3)):
                # (batch, 1, 28, 28)
                logit = self.forward(x)
                probability = torch.sigmoid(logit)
                sample = torch.bernoulli(probability)
                x[:, :, i, j] = sample[:, :, i, j]
        
        return x
model = PixelCNN().cuda()
optimizer = torch.optim.Adam(model.parameters(), 1e-3)
print('done')
from IPython import display


step = 0
while(True):
    for batch_index, batch in enumerate(train_loader):
        # x : image (B, 1, 28, 28,)
        # t : class (B)
        x, t = batch
        
        x = x.cuda()
        x[x >= 0.5] = 1.
        x[x < 0.5] = 0.

        # Backpropagation이 되도록 준비합니다.       
        model.train()    
        model.zero_grad()
        y = model(x)
        loss = nn.BCEWithLogitsLoss()(y, x)
        loss.backward()
        optimizer.step()
        
        if step % 100 == 0:
            print(step, loss.item())
        
        if step % 1000 == 0:
            display.clear_output()
            
            samples = torch.zeros(10, 1, 28, 28).cuda()
            model.eval()
            with torch.no_grad():
                samples = model.inference(x)
            _samples = samples.data.cpu().numpy()
            
            plt.figure(figsize=[18, 3])
            for i in range(10):
                plt.subplot(1, 10, i+1)
                plt.imshow(_samples[i, 0], cmap='Greys')
                plt.xticks([])
                plt.yticks([])
            plt.show()
        
        step += 1<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11791762492159732, 'fraction_numerical': 0.019234789880827934, 'mean_word_length': 2.0374603174603174, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '307901', 'license': 'MIT', 'n_tokens_mistral': 1648, 'n_tokens_neox': 1579, 'n_words': 326, 'path': 'Autoregressive/PixelCNN.ipynb', 'repo_name': 'scpark20/DGM-colab', 'size': '4783'}