<jupyter_start><jupyter_text>Import libraries<jupyter_code>import torch
import torchvision
from torch import nn
from torch import optim
import torch.nn.functional as F
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.utils import save_image
from torchvision.datasets import MNIST
from torchvision import datasets
import os<jupyter_output><empty_output><jupyter_text>Train Autoencoder<jupyter_code>if not os.path.exists('./vae_img'):
    os.mkdir('./vae_img')


def to_img(x):
    x = x.clamp(0, 1)
    x = x.view(x.size(0), 1, 28, 28)
    return x


num_epochs = 100
batch_size = 128
learning_rate = 1e-3

img_transform = transforms.Compose([
    transforms.ToTensor()
    # transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

dataset = MNIST('./data', transform=img_transform, download=True)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)


class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()

        self.fc1 = nn.Linear(784, 400)
        self.fc21 = nn.Linear(400, 20)
        self.fc22 = nn.Linear(400, 20)
        self.fc3 = nn.Linear(20, 400)
        self.fc4 = nn.Linear(400, 784)

    def encode(self, x):
        h1 = F.relu(self.fc1(x))
        return self.fc21(h1), self.fc22(h1)

    def reparametrize(self, mu, logvar):
        std = logvar.mul(0.5).exp_()
        if torch.cuda.is_available():
            eps = torch.cuda.FloatTensor(std.size()).normal_()
        else:
            eps = torch.FloatTensor(std.size()).normal_()
        eps = Variable(eps)
        return eps.mul(std).add_(mu)

    def decode(self, z):
        h3 = F.relu(self.fc3(z))
        return F.sigmoid(self.fc4(h3))

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparametrize(mu, logvar)
        return self.decode(z), mu, logvar


model = VAE()
if torch.cuda.is_available():
    model.cuda()

reconstruction_function = nn.MSELoss(size_average=False)


def loss_function(recon_x, x, mu, logvar):
    """
    recon_x: generating images
    x: origin images
    mu: latent mean
    logvar: latent log variance
    """
    BCE = reconstruction_function(recon_x, x)  # mse loss
    # loss = 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
    KLD_element = mu.pow(2).add_(logvar.exp()).mul_(-1).add_(1).add_(logvar)
    KLD = torch.sum(KLD_element).mul_(-0.5)
    # KL divergence
    return BCE + KLD


optimizer = optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(num_epochs):
    model.train()
    train_loss = 0
    for batch_idx, data in enumerate(dataloader):
        img, _ = data
        img = img.view(img.size(0), -1)
        img = Variable(img)
        if torch.cuda.is_available():
            img = img.cuda()
        optimizer.zero_grad()
        recon_batch, mu, logvar = model(img)
        loss = loss_function(recon_batch, img, mu, logvar)
        loss.backward()
        train_loss += loss.data
        optimizer.step()
        if batch_idx % 100 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch,
                batch_idx * len(img),
                len(dataloader.dataset), 100. * batch_idx / len(dataloader),
                loss.data / len(img)))

    print('====> Epoch: {} Average loss: {:.4f}'.format(
        epoch, train_loss / len(dataloader.dataset)))
    if epoch % 10 == 0:
        save = to_img(recon_batch.cpu().data)
        save_image(save, './vae_img/image_{}.png'.format(epoch))

torch.save(model.state_dict(), './vae.pth')<jupyter_output>Train Epoch: 0 [0/60000 (0%)]	Loss: 183.676865
Train Epoch: 0 [12800/60000 (21%)]	Loss: 45.187099
Train Epoch: 0 [25600/60000 (43%)]	Loss: 42.488605
Train Epoch: 0 [38400/60000 (64%)]	Loss: 38.989422
Train Epoch: 0 [51200/60000 (85%)]	Loss: 37.602478
====> Epoch: 0 Average loss: 45.2671
Train Epoch: 1 [0/60000 (0%)]	Loss: 37.402729
Train Epoch: 1 [12800/60000 (21%)]	Loss: 34.094231
Train Epoch: 1 [25600/60000 (43%)]	Loss: 36.393234
Train Epoch: 1 [38400/60000 (64%)]	Loss: 35.329041
Train Epoch: 1 [51200/60000 (85%)]	Loss: 33.767498
====> Epoch: 1 Average loss: 34.8616
Train Epoch: 2 [0/60000 (0%)]	Loss: 33.174458
Train Epoch: 2 [12800/60000 (21%)]	Loss: 33.061054
Train Epoch: 2 [25600/60000 (43%)]	Loss: 34.152504
Train Epoch: 2 [38400/60000 (64%)]	Loss: 33.052860
Train Epoch: 2 [51200/60000 (85%)]	Loss: 34.183891
====> Epoch: 2 Average loss: 33.0455
Train Epoch: 3 [0/60000 (0%)]	Loss: 33.032677
Train Epoch: 3 [12800/60000 (21%)]	Loss: 31.276497
Train Epoch: 3 [25600/60000 (43%)]	Loss: [...]<jupyter_text>Use trained model to predict images<jupyter_code>IMG_PATH = "D:\\University-Work\CPS040\Main\CPS040-Thesis\Main\Autoencoder\Simple Autoencoder 2\\vae_img\\"
min_img_size = 32
trans = transforms.Compose([transforms.Resize(min_img_size),
                                         transforms.CenterCrop(32),
                                         transforms.ToTensor(),
                                         transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                                              std=[0.229, 0.224, 0.225])])

# # Picture dataset.
# classify_dataset = datasets.ImageFolder(root=IMG_PATH, transform=trans)
# # Create custom random sampler class to iter over dataloader.
# classify_loader = DataLoader(dataset=classify_dataset, batch_size=1, shuffle=True, num_workers=5)

model = torch.load('./vae.pth')

# if cuda is available, move the model to the GPU
# if torch.cuda.is_available():
#     model.cuda()

def to_img(x):
    x = x.clamp(0, 1)
    x = x.view(x.size(0), 1, 28, 28)
    return x
img_transform = transforms.Compose([
    transforms.ToTensor()
    # transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

dataset = MNIST('./data', transform=img_transform, download=True)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# for batch_idx, data in enumerate(dataloader):
#     data

img, _ = data
img = img.view(img.size(0), -1)
img = Variable(img)
if torch.cuda.is_available():
    img = img.cuda()
    
model = torch.load('./vae.pth')
model(data)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14948119325551232, 'fraction_numerical': 0.08690012970168612, 'mean_word_length': 2.936821952776005, 'pattern_counts': {'":': 0, '<': 12, '<?xml version=': 0, '>': 16, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '184306', 'license': 'MIT', 'n_tokens_mistral': 2666, 'n_tokens_neox': 2247, 'n_words': 511, 'path': 'Main/Autoencoder/Simple Autoencoder 2/PyTorch Autoencoder.ipynb', 'repo_name': 'MalcolmGomes/CPS040-Thesis', 'size': '6168'}