Upload Untitled-1.py

Untitled-1.py

@@ -0,0 +1,450 @@
+# %%
+import pandas
+import torch
+import matplotlib.pyplot as plt
+from torch import nn
+
+# %%
+class Patches_to_Embedding(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.flatten = nn.Flatten(2,-1)
+        self.conv = nn.Conv2d(in_channels=3,out_channels=384,kernel_size=8,stride=8)
+        self.cls_token = nn.Parameter(torch.randn(1,1,384))
+        self.position_embedding = nn.Parameter(torch.randn(1,65,384))
+
+    def forward(self,x :torch.Tensor):
+        batch_size = x.shape[0]
+        x = self.conv(x)
+        x = self.flatten(x)
+        x = x.transpose(1,2)
+        cls = self.cls_token.expand(batch_size,-1,-1)
+        x = torch.cat((cls,x),dim=1)
+        x = x + self.position_embedding
+        return x
+
+# %%
+class EncoderBlock(nn.Module):
+    def __init__(self,dropout):
+        super().__init__()
+        self.dropout = dropout
+        self.Layer_Norm1 = nn.LayerNorm(384)
+        self.Layer_Norm2 = nn.LayerNorm(384)
+        self.multi_head_attention = nn.MultiheadAttention(384,6,batch_first=True,dropout=dropout)
+        self.MLP = nn.Sequential(
+            nn.Linear(384,1536),
+            nn.GELU(),
+            nn.Dropout(self.dropout),
+            nn.Linear(1536,384),
+            nn.Dropout(self.dropout)
+        )
+
+    def forward(self,x):
+        output = self.Layer_Norm1(x)
+        output,_ = self.multi_head_attention(output,output,output)
+        x = x + output
+        output = self.Layer_Norm2(x)
+        output = self.MLP(output)
+        x = output + x
+        return x
+
+# %%
+class Encoder(nn.Module):
+    def __init__(self,dropout,num_layers):
+        super().__init__()
+        self.layers = nn.ModuleList([
+            EncoderBlock(dropout)
+            for _ in range(num_layers)
+        ])
+        self.norm = nn.LayerNorm(384)
+
+    def forward(self,x:torch.Tensor) ->torch.Tensor:
+        for layer in self.layers:
+            x = layer(x)
+        x = self.norm(x)
+        return x
+
+# %%
+class ViT(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.patch_embedding = Patches_to_Embedding()
+        self.encoder = Encoder(0.2,8)
+        self.head = nn.Linear(384,39)
+
+    def forward(self,x):
+        x = self.patch_embedding(x)
+        x = self.encoder(x)
+        cls = x[:,0]
+        logits = self.head(cls)
+        return logits
+
+# %%
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader, Subset
+import torch
+
+DATA_DIR = "/home/ujwal/Documents/Pytorch/GitHub/Vision_Transformer/Plant_leave_diseases_dataset_without_augmentation"
+
+train_transform = transforms.Compose([
+    transforms.RandomResizedCrop(
+        64,
+        scale=(0.8, 1.0),
+        ratio=(0.9, 1.1)
+    ),
+
+    transforms.RandomHorizontalFlip(p=0.5),
+    transforms.RandomVerticalFlip(p=0.5),
+
+    transforms.RandomRotation(15),
+
+    transforms.ColorJitter(
+        brightness=0.2,
+        contrast=0.2,
+        saturation=0.2,
+        hue=0.05
+    ),
+
+    transforms.ToTensor(),
+
+    transforms.Normalize(
+        mean=[0.485, 0.456, 0.406],
+        std=[0.229, 0.224, 0.225]
+    ),
+
+    transforms.RandomErasing(
+        p=0.25,
+        scale=(0.02, 0.2),
+        ratio=(0.3, 3.3)
+    )
+])
+
+test_transform = transforms.Compose([
+    transforms.Resize((64, 64)),
+
+    transforms.ToTensor(),
+
+    transforms.Normalize(
+        mean=[0.485, 0.456, 0.406],
+        std=[0.229, 0.224, 0.225]
+    )
+])
+
+# %%
+full_dataset = datasets.ImageFolder(DATA_DIR)
+
+train_dataset_full = datasets.ImageFolder(
+    DATA_DIR,
+    transform=train_transform
+)
+
+test_dataset_full = datasets.ImageFolder(
+    DATA_DIR,
+    transform=test_transform
+)
+
+generator = torch.Generator().manual_seed(42)
+
+train_size = int(0.8 * len(full_dataset))
+test_size = len(full_dataset) - train_size
+
+train_subset, test_subset = torch.utils.data.random_split(
+    range(len(full_dataset)),
+    [train_size, test_size],
+    generator=generator
+)
+
+train_indices = train_subset.indices
+test_indices = test_subset.indices
+
+train_dataset = Subset(
+    train_dataset_full,
+    train_indices
+)
+
+test_dataset = Subset(
+    test_dataset_full,
+    test_indices
+)
+
+print("Train:", len(train_dataset))
+print("Test:", len(test_dataset))
+
+# %%
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device
+
+# %%
+import os
+
+model = ViT().to(device)
+model = torch.compile(model)
+
+checkpoint_path = "checkpoint.pth"
+
+if os.path.exists(checkpoint_path):
+
+    checkpoint = torch.load(
+        checkpoint_path,
+        map_location=device
+    )
+
+    model.load_state_dict(checkpoint["model"])
+
+    best_acc = checkpoint["best_acc"]
+    start_epoch = checkpoint["epoch"] + 1
+
+    print(f"Resuming from epoch {start_epoch}")
+    print(f"Best validation accuracy: {best_acc:.2f}%")
+
+else:
+
+    best_acc = 0.0
+    start_epoch = 0
+
+    print("No checkpoint found.")
+    print("Starting training from scratch.")
+
+# %%
+criterion = nn.CrossEntropyLoss(
+    label_smoothing=0.1
+)
+
+optimizer = torch.optim.AdamW(
+    model.parameters(),
+    lr=1e-4,
+    weight_decay=0.03
+)
+
+scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+    optimizer,
+    T_max=50
+)
+
+# %%
+class AugmentedDataset(torch.utils.data.Dataset):
+
+    def __init__(self, dataset):
+        self.dataset = dataset
+
+    def __len__(self):
+        return len(self.dataset) * 4
+
+    def __getitem__(self, index):
+
+        original_index = index // 4
+        version = index % 4
+
+        image, label = self.dataset[original_index]
+
+        if version == 0:
+            return image, label
+
+        elif version == 1:
+            return torch.flip(image, dims=[2]), label
+
+        elif version == 2:
+            return torch.flip(image, dims=[1]), label
+
+        else:
+            return torch.flip(image, dims=[1, 2]), label
+
+# %%
+train_dataset = AugmentedDataset(train_dataset)
+
+# %%
+print(len(train_dataset))
+
+# %%
+from torch.utils.data import DataLoader
+
+train_loader = DataLoader(
+    train_dataset,
+    batch_size=32,
+    shuffle=True,
+    num_workers=4,
+    pin_memory=True,
+    persistent_workers=True,
+    prefetch_factor=2
+)
+
+test_loader = DataLoader(
+    test_dataset,
+    batch_size=32,
+    shuffle=False,
+    num_workers=4,
+    pin_memory=True,
+    persistent_workers=True,
+    prefetch_factor=2
+)
+
+# %%
+from tqdm.auto import tqdm
+
+scaler = torch.amp.GradScaler("cuda")
+total_epochs = 50
+
+for epoch in range(start_epoch, total_epochs):
+
+    model.train()
+
+    running_loss = 0
+    correct = 0
+    total = 0
+
+    progress_bar = tqdm(
+        train_loader,
+        desc=f"Epoch [{epoch+1}/{total_epochs}]",
+        leave=True
+    )
+
+    for images, labels in progress_bar:
+
+        images = images.to(device, non_blocking=True)
+        labels = labels.to(device, non_blocking=True)
+
+        optimizer.zero_grad()
+
+        with torch.amp.autocast(device_type="cuda"):
+            logits = model(images)
+            loss = criterion(logits, labels)
+
+        scaler.scale(loss).backward()
+        scaler.step(optimizer)
+        scaler.update()
+
+        running_loss += loss.item()
+
+        predictions = logits.argmax(dim=1)
+
+        correct += (predictions == labels).sum().item()
+        total += labels.size(0)
+
+        progress_bar.set_postfix(
+            loss=running_loss / len(progress_bar),
+            accuracy=100 * correct / total
+        )
+
+    train_loss = running_loss / len(train_loader)
+    train_accuracy = 100 * correct / total
+
+    model.eval()
+
+    test_loss = 0
+    correct = 0
+    total = 0
+
+    with torch.no_grad():
+
+        progress_bar = tqdm(
+            test_loader,
+            desc="Testing",
+            leave=False
+        )
+
+        for images, labels in progress_bar:
+
+            images = images.to(device, non_blocking=True)
+            labels = labels.to(device, non_blocking=True)
+
+            with torch.amp.autocast(device_type="cuda"):
+                logits = model(images)
+                loss = criterion(logits, labels)
+
+            test_loss += loss.item()
+
+            predictions = logits.argmax(dim=1)
+
+            correct += (predictions == labels).sum().item()
+            total += labels.size(0)
+
+    test_loss /= len(test_loader)
+    test_accuracy = 100 * correct / total
+
+    print(
+        f"Epoch {epoch+1}: "
+        f"Train Loss={train_loss:.4f}, "
+        f"Train Accuracy={train_accuracy:.2f}%, "
+        f"Test Loss={test_loss:.4f}, "
+        f"Test Accuracy={test_accuracy:.2f}%"
+    )
+
+    scheduler.step()
+
+    if test_accuracy > best_acc:
+        best_acc = test_accuracy
+
+        torch.save({
+            "epoch": epoch,
+            "model": model.state_dict(),
+            "optimizer": optimizer.state_dict(),
+            "scheduler": scheduler.state_dict(),
+            "best_acc": best_acc
+        }, "checkpoint.pth")
+
+# %%
+test_loader = DataLoader(
+    test_dataset,
+    batch_size=32,
+    shuffle=False,
+    num_workers=4,
+    pin_memory=True
+)
+
+# %%
+import random
+import matplotlib.pyplot as plt
+
+def predict_test_image(index=None):
+
+    if index is None:
+        index = random.randrange(len(test_dataset))
+
+    image_tensor, actual_label = test_dataset[index]
+
+    original_image, _ = full_dataset[test_indices[index]]
+
+    image_input = image_tensor.unsqueeze(0).to(device)
+
+    model.eval()
+
+    with torch.no_grad():
+        logits = model(image_input)
+
+        probabilities = torch.softmax(logits, dim=1)
+
+        predicted_label = logits.argmax(dim=1).item()
+
+        confidence = probabilities[0, predicted_label].item()
+
+    actual_class = full_dataset.classes[actual_label]
+    predicted_class = full_dataset.classes[predicted_label]
+
+    # Show image
+    plt.figure(figsize=(6, 6))
+    plt.imshow(original_image)
+    plt.axis("off")
+
+    plt.title(
+        f"Actual: {actual_class}\n"
+        f"Predicted: {predicted_class}\n"
+        f"Confidence: {confidence * 100:.2f}%"
+    )
+
+    plt.show()
+
+    print("Test index:", index)
+    print("Actual:", actual_class)
+    print("Predicted:", predicted_class)
+    print(f"Confidence: {confidence * 100:.2f}%")
+
+# %%
+predict_test_image()
+
+# %%
+predict_test_image(108)
+
+# %%
+
+
+