<jupyter_start><jupyter_text><jupyter_code>import torch as t
import torchvision
from tqdm import tqdm
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
from matplotlib.colors import Colormap
train_data = torchvision.datasets.FashionMNIST("data", train=True, transform=transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]), download=True)
train_loader = t.utils.data.DataLoader(train_data, batch_size = 128, shuffle=True, num_workers=2)
!pip install einops
import einops
plt.imshow(train_data[0][0].squeeze(), plt.cm.binary)
ngc = 64
latent_size = 100
num_classes = 10

class Generator(t.nn.Module):
  def __init__(self, 
               latent_size=latent_size, 
               num_classes = num_classes, 
               emb_size = 96, 
               nhead=4,
               L1 = 4,
               L2 = 4,
               img_size = 28,
               in_channels=1):
    super().__init__()
    self.emb_size = emb_size
    self.img_size = img_size
    # self.latent_size = latent_size
    self.linear1 = t.nn.Linear(latent_size, emb_size*(img_size//4)**2)
    self.block1 = t.nn.Sequential(
        *[t.nn.TransformerEncoderLayer(
        emb_size,
        nhead=nhead, 
        dim_feedforward=emb_size*4, 
        activation="gelu",
        norm_first=True,
        batch_first=True, 
        dropout=0., 
        layer_norm_eps=1e-12
      ) for _ in range(L1)]
    )
    self.pixel1 = t.nn.PixelShuffle(2)
    self.block2 = t.nn.Sequential(
        *[t.nn.TransformerEncoderLayer(
        emb_size//4,
        nhead=nhead, 
        dim_feedforward=emb_size, 
        activation="gelu",
        norm_first=True,
        batch_first=True, 
        dropout=0., 
        layer_norm_eps=1e-12
      ) for _ in range(L2)]
    )
    self.pixel2 = t.nn.PixelShuffle(2)
    self.project = t.nn.Linear(emb_size//16, in_channels)
    self.tanh = t.nn.Tanh()

  def forward(self, z):
    b = z.shape[0]
    emb = self.linear1(z).view((b, (self.img_size//4)**2, self.emb_size))
    emb = self.block1(emb)
    emb_to_pixel = "b (h w) c -> b c h w"
    pixel_to_emb = "b c h w -> b (h w) c"

    pixels = einops.rearrange(emb, emb_to_pixel, h = self.img_size//4)
    pixels = self.pixel1(pixels)
    emb = einops.rearrange(pixels, pixel_to_emb, h = self.img_size//2)
    emb = self.block2(emb)
    pixels = einops.rearrange(emb, emb_to_pixel, h = self.img_size//2)
    pixels = self.pixel2(pixels)
    emb = einops.rearrange(pixels, pixel_to_emb, h = self.img_size)
    emb = self.project(emb)
    pixels = einops.rearrange(emb, emb_to_pixel, h = self.img_size)
    return self.tanh(pixels)

# class Generator(t.nn.Module):
#   def __init__(self, latent_size=latent_size, num_classes = num_classes):
#     super().__init__()
#     self.layers = t.nn.Sequential(
#         t.nn.ConvTranspose2d(latent_size, ngc*4, 4, 1, 0),
#         t.nn.BatchNorm2d(ngc*4),
#         t.nn.ReLU(),
#         t.nn.ConvTranspose2d(ngc*4, ngc*2, 4, 2, 1),
#         t.nn.BatchNorm2d(ngc*2),
#         t.nn.ReLU(),
#         t.nn.ConvTranspose2d(ngc*2, ngc*1, 4, 2, 1),
#         t.nn.BatchNorm2d(ngc),
#         t.nn.ReLU(),
#         t.nn.ConvTranspose2d(ngc*1, 1, 4, 2, 3),
#         t.nn.Tanh()
#     )
#   def forward(self, z):
#     return self.layers(z.view(-1, latent_size, 1, 1))

class PatchEmbedding(t.nn.Module):
  def __init__(self, in_channels, patch_size, emb_size, imgsize):
    super().__init__()
    self.cls_token = t.nn.Parameter(t.randn(1, 1, emb_size)) # b, n, emb_size. Add to the list of module params
    self.n = (imgsize//patch_size)**2
    self.emb_size = emb_size
    self.position_embeddings = t.nn.Parameter(t.randn(1, self.n + 1, emb_size)) # +1 for cls
    self.projection = t.nn.Conv2d(in_channels, emb_size, patch_size, patch_size, 0)

  def forward(self, inputs):
    token_embeddings = einops.rearrange(self.projection(inputs), 'b c h w -> b (h w) c')
    b = inputs.shape[0]
    cls_embeddings = self.cls_token.repeat((b, 1, 1))
    # print(cls_embeddings.shape, token_embeddings.shape, self.position_embeddings.shape)
    return t.cat((cls_embeddings, token_embeddings), dim = 1) + self.position_embeddings

class ClassificationHead(t.nn.Module):
  def __init__(self, num_classes, emb_size):
    super().__init__()
    self.ln = t.nn.LayerNorm((emb_size,), eps = 1e-12)
    self.layer = t.nn.Linear(emb_size, num_classes)

  def forward(self, x):
    return self.layer(self.ln(x)[:,0,:])


class ViT_New(t.nn.Module):
  def __init__(self, in_channels: int = 1, patch_size: int = 4, emb_size: int = 96, nhead=6, imgsize=28, num_classes=1, L = 6, from_hugging = False):
    super().__init__()
    self.patch_emb = PatchEmbedding(in_channels, patch_size, emb_size, imgsize)
    # self.from_hugging = from_hugging
    # if from_hugging:
    #   config = ViTConfig(hidden_size=emb_size, num_hidden_layers = L, num_attention_heads = nhead, intermediate_size=emb_size*4, patch_size=patch_size, image_size=imgsize, encoder_stride=patch_size)
    #   self.encoder = ViTForImageClassification(config).vit.encoder
    # else:
    self.encoder = t.nn.Sequential(*[
      t.nn.TransformerEncoderLayer(
        emb_size,
        nhead=nhead, 
        dim_feedforward=emb_size*4, 
        activation="gelu",
        norm_first=True,
        batch_first=True, 
        dropout=0., 
        layer_norm_eps=1e-12
      ) for _ in range(L)])
    self.classifier = ClassificationHead(num_classes, emb_size=emb_size)

  def forward(self, x):
    emb = self.patch_emb(x)
    encoding = self.encoder(emb)
    return t.sigmoid(self.classifier(encoding)).flatten()
    



lr = 0.0002
beta1 = 0.
gen = Generator().cuda()
disc = ViT_New().cuda()
optim_g = t.optim.Adam(gen.parameters(), lr=lr, betas=(beta1, 0.999))
optim_d = t.optim.Adam(disc.parameters(), lr=lr, betas=(beta1, 0.999))

gz = gen(t.randn(1, 100).cuda())
print(gz.shape)
disc(gz)
criterion = t.nn.BCELoss()
# training loop
num_epochs = 5
real_label = 1.
fake_label = 0.

d_losses = []
g_losses = []
for epoch in range(num_epochs):
  for step, (x, y) in enumerate(tqdm(train_loader)):
    # lab = t.nn.functional.one_hot(y.cuda(), num_classes = num_classes).float()
    # update G
    gen.zero_grad()
    disc.zero_grad()
    b = x.shape[0]
    with t.no_grad():
      z = t.randn(b, latent_size).cuda()
    gz = gen(z)
    dgz = disc(gz)
    # print(dgz.shape, b)
    g_loss = criterion(dgz, t.full((b,), real_label).cuda())
    g_losses.append(g_loss.item())
    g_loss.backward()
    optim_g.step()
    # update D
    disc.zero_grad()
    dx = disc(x.cuda())
    dgz = disc(gz.detach())
    d_loss = criterion(dgz, t.full((b,), fake_label).cuda()) + criterion(dx, t.full((b,), real_label).cuda())
    d_losses.append(d_loss.item())
    d_loss.backward()
    optim_d.step()
    n_steps = 20
    if step % n_steps == 0:
      L_g = t.tensor(g_losses[-n_steps:]).mean()
      L_d = t.tensor(d_losses[-n_steps:]).mean()
      print(f'L_g: {L_g:.4f}, L_d: {L_d:.4f}, dx: {dx.detach().mean().item():.4f}, dgz: {dgz.detach().mean().item():.4f}')
  



with t.no_grad():
  z = t.randn(50, latent_size).cuda()
  # lab = t.nn.functional.one_hot(t.arange(50).cuda() % 10, num_classes = num_classes).float()

  ims = gen(z).cpu()
  _, axs = plt.subplots(5, 10)
  axs = axs.flatten()
  for im, ax in zip(ims, axs):
    ax.axes.xaxis.set_visible(False)
    ax.axes.yaxis.set_visible(False)

    ax.imshow(im.squeeze(), cmap="gray_r")
    
  
  plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15151920760273055, 'fraction_numerical': 0.02195154597778075, 'mean_word_length': 3.0697167755991286, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '293245', 'license': 'MIT', 'n_tokens_mistral': 3052, 'n_tokens_neox': 2913, 'n_words': 514, 'path': 'TransformerGAN.ipynb', 'repo_name': 'The20thDuck/Neuro-140-Project', 'size': '7471'}