<jupyter_start><jupyter_text>CNN VAE for SRL for DIY Self driving carIn this notebook you will learn the CNN VAE(beta) model. The result model is used for state representation in reinforcement learning.First collection training data. you can use notebooks\utility\data_collection.ipynb Collect images of the course while driving the car on the course. Collect 1k to 10k images. Adjust the number of data collected according to the size of the course. When running the course, run in the center of the course, the side of the side line, zigzag running, etc. During the trial during reinforcement learning, you do not know how to run on the course. Collect data so that the course can be represented in the event of an error. Mount google driveYou upload zip file that contain training data. The zip file copy from googledrive. Set zip file name to DATASET_ZIP.<jupyter_code>from google.colab import drive 
drive.mount('/content/drive')

DATASET_ZIP = 'jetracer_dataset.zip'
DATASET_DIR = 'dataset'<jupyter_output><empty_output><jupyter_text>Copy from google driveCopy training data and unzip.<jupyter_code>!rm -rf dataset_root
!cp '/content/drive/My Drive/$DATASET_ZIP' ./
!unzip -q $DATASET_ZIP

!mkdir dataset_root
!mv $DATASET_DIR './dataset_root'<jupyter_output><empty_output><jupyter_text>Resize dataset<jupyter_code>import PIL
import glob,os

files = glob.glob(os.path.join('/content/dataset_root', DATASET_DIR, '*.jpg'))
for f in files:
  try:
    image = PIL.Image.open(f)
  except OSError:
    print('Delete' + f)
    !rm -rf f
  image = image.resize((160,120))
  image.crop((0, 40, 160, 120)).save(f, quality=95)<jupyter_output><empty_output><jupyter_text>Import module<jupyter_code>import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torchvision import datasets
from torchvision import transforms
from torchvision.utils import save_image
from IPython.display import Image
from IPython.core.display import Image, display

%load_ext autoreload
%autoreload 2<jupyter_output><empty_output><jupyter_text>Load GPU device<jupyter_code>device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')<jupyter_output><empty_output><jupyter_text>Load dataset<jupyter_code>bs = 64
dataset = datasets.ImageFolder(root='./dataset_root', transform=transforms.Compose([
    transforms.ToTensor(),
]))
dataloader = torch.utils.data.DataLoader(dataset, batch_size=bs, shuffle=True)
len(dataset.imgs), len(dataloader)<jupyter_output><empty_output><jupyter_text><jupyter_code>fixed_x, _ = next(iter(dataloader))
save_image(fixed_x, 'real_image.png')
Image('real_image.png')<jupyter_output><empty_output><jupyter_text>Define VAE Network<jupyter_code>class Flatten(nn.Module):
    def forward(self, input):
        return input.view(input.size(0), -1)

class UnFlatten(nn.Module):
    def forward(self, input, size=256):
        return input.view(input.size(0), size, 3, 8)


class VAE(nn.Module):
    def __init__(self, image_channels=3, h_dim=6144, z_dim=32):
        super(VAE, self).__init__()
        self.z_dim = z_dim
        self.encoder = nn.Sequential(
            nn.Conv2d(image_channels, 32, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(128, 256, kernel_size=4, stride=2),
            nn.ReLU(),
            Flatten()
        )
        
        self.fc1 = nn.Linear(h_dim, z_dim)
        self.fc2 = nn.Linear(h_dim, z_dim)
        self.fc3 = nn.Linear(z_dim, h_dim)
        
        self.decoder = nn.Sequential(
            UnFlatten(),
            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.ConvTranspose2d(64, 32, kernel_size=5, stride=2),
            nn.ReLU(),
            nn.ConvTranspose2d(32, image_channels, kernel_size=4, stride=2),
            nn.Sigmoid(),
        )
        
    def reparameterize(self, mu, logvar):
        std = logvar.mul(0.5).exp_()
        esp = torch.randn(*mu.size()).to(device)
        z = mu + std * esp
        return z
    
    def bottleneck(self, h):
        mu, logvar = self.fc1(h), F.softplus(self.fc2(h))
        z = self.reparameterize(mu, logvar)
        return z, mu, logvar

    def encode(self, x):
        h = self.encoder(x)
        z, mu, logvar = self.bottleneck(h)
        return z, mu, logvar

    def decode(self, z):
        z = self.fc3(z)
        z = self.decoder(z)
        return z

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

    def loss_fn(self, images, reconst, mean, logvar):
        KL = -0.5 * torch.sum((1 + logvar - mean.pow(2) - logvar.exp()), dim=0)
        KL = torch.mean(KL)
        reconstruction = F.binary_cross_entropy(reconst.view(-1,38400), images.view(-1, 38400), reduction='sum') #size_average=False)
        return reconstruction + 5.0 * KL<jupyter_output><empty_output><jupyter_text>Prepare TrainingCreate VAE model and initialize optimizer.<jupyter_code>from torchsummary import summary
VARIANTS_SIZE = 32
image_channels = fixed_x.size(1)
vae = VAE(image_channels=image_channels, z_dim=VARIANTS_SIZE ).to(device)
optimizer = torch.optim.Adam(vae.parameters(), lr=1e-3)
summary(vae, (3, 80, 160))<jupyter_output><empty_output><jupyter_text>Tensorboard<jupyter_code>%load_ext tensorboard
%tensorboard --logdir ./runs<jupyter_output><empty_output><jupyter_text>Start training<jupyter_code>from torch.utils.tensorboard import SummaryWriter
import numpy as np
epochs = 50
writer = SummaryWriter()

vae.train()
for epoch in range(epochs):
    losses = []
    grid = None
    for idx, (images, _) in enumerate(dataloader):
        images = images.to(device)
        optimizer.zero_grad()
        recon_images, mu, logvar = vae(images)
        loss = vae.loss_fn(images, recon_images, mu, logvar)
        loss.backward()
        optimizer.step()
        losses.append(loss.cpu().detach().numpy())
        grid = torchvision.utils.make_grid(recon_images)
    writer.add_image('Image/reconst', grid, epoch)
    writer.add_scalar('Loss/train',np.average(losses), epoch)
    print("EPOCH: {} loss: {}".format(epoch+1, np.average(losses)))

torch.save(vae.state_dict(), 'vae.torch')<jupyter_output><empty_output><jupyter_text>CleanupCopy trained model file to GoogleDrive.<jupyter_code>!cp vae.torch '/content/drive/My Drive/vae.torch'<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12020033388981637, 'fraction_numerical': 0.020792229473364697, 'mean_word_length': 3.680397727272727, 'pattern_counts': {'":': 0, '<': 49, '<?xml version=': 0, '>': 49, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '12', 'dir': 'jupyter-structured-clean-dedup', 'id': '96652', 'license': 'MIT', 'n_tokens_mistral': 2343, 'n_tokens_neox': 2241, 'n_words': 519, 'path': 'notebooks/colabo/VAE_CNN.ipynb', 'repo_name': 'ko4ro/rl4aicar', 'size': '6589'}