| import kagglehub
|
| import cv2
|
| import os
|
| from IPython.display import clear_output
|
| import torch
|
| import numpy as np
|
| from torch.utils.data import Dataset, DataLoader
|
| import torch.nn as nn
|
| import matplotlib.pyplot as plt
|
| from autoencoder import Encoder, Decoder
|
| from trainer import Trainer
|
| from objectives import Discriminator, vgg_builder
|
|
|
|
|
| image_shape = 256
|
| emb_dim = 768
|
| patch_size = 16
|
|
|
| image_path = kagglehub.dataset_download("awsaf49/coco-2017-dataset")
|
| data = []
|
| for dirpath, _, filenames in os.walk(image_path):
|
| for filename in filenames:
|
| if filename.endswith("jpg"):
|
| name = os.path.join(dirpath, filename)
|
| img = cv2.imread(name)
|
| img = cv2.resize(img, (image_shape,image_shape))
|
| img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
| img = img.astype(np.float32) / 127.5 - 1.0
|
| img = torch.tensor(img).permute(2,0,1)
|
| data.append(img)
|
| clear_output(wait=1)
|
| print(f"{len(data)/1670:.2f}%")
|
| print(len(data))
|
|
|
| class CustomDataset(Dataset):
|
| def __init__(self, data):
|
| self.indices = np.arange(len(data))
|
| np.random.shuffle(self.indices)
|
| self.data = data
|
|
|
| def __len__(self):
|
| return len(self.indices)
|
|
|
| def __getitem__(self, idx):
|
| return torch.tensor(self.data[self.indices[idx]], dtype=torch.float32)
|
|
|
|
|
| plt.imshow(CustomDataset(data)[0].permute(1,2,0)/2+0.5)
|
|
|
| encoder = Encoder(latent_dim=16)
|
| decoder = Decoder(latent_dim=16)
|
| D = Discriminator((3,256,256))
|
|
|
| vgg = vgg_builder()
|
| for param in vgg.parameters():
|
| param.requires_grad = False
|
| vgg.eval()
|
|
|
| print(f"encoder: {sum(p.numel() for p in encoder.parameters())/(262144):.3f}MB")
|
| print(f"decoder: {sum(p.numel() for p in decoder.parameters())/(262144):.3f}MB")
|
| print(f"Discriminator: {sum(p.numel() for p in D.parameters())/(262144):.3f}MB")
|
| print(f"VGG: {sum(p.numel() for p in vgg.parameters())/(262144):.3f}MB")
|
|
|
| batch_size = 16
|
| dataset = CustomDataset(data)
|
| loader = DataLoader(dataset,
|
| batch_size=batch_size,
|
| shuffle=True,
|
| num_workers=8,
|
| pin_memory=True)
|
| epochs = 5
|
| trainer = Trainer(encoder, decoder, D, vgg, ["mse", "gan", "vgg", "KL"], len(loader) if "loader" in locals() else 0, isViT=1)
|
| for epoch in range(1, epochs):
|
| index = 0
|
| for i, x in enumerate(loader):
|
| trainer.train_step(x, freeze_disc=0, with_mse=1, freeze_ae=0)
|
| trainer.update_epoch()
|
|
|
| torch.save(encoder.state_dict(), "encoder16.pt")
|
| torch.save(decoder.state_dict(), "decoder16.pt")
|
| torch.save(D.state_dict(), "discriminator16.pt") |
