import torch import zipfile import torch.nn as nn import torch.optim as optim from torchvision import models, transforms from torch.utils.data import Dataset, DataLoader, ConcatDataset, Subset import numpy as np import random import os import tarfile import pickle from sklearn.metrics import accuracy_score import copy import pandas as pd  # --------------------- # 1. Utility Functions # ---------------------  def set_seed(seed=42):     """Set random seeds for reproducibility."""     random.seed(seed)     np.random.seed(seed)     torch.manual_seed(seed)     torch.cuda.manual_seed_all(seed)  set_seed()  # Device configuration device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}")  # --------------------- # 2. Custom Dataset Class # ---------------------  class TarPTHDataset(Dataset):     """     Custom Dataset to load data from .tar.pth files.     Assumes each .tar.pth file contains a dictionary with 'data' and 'labels' keys.     - 'data': Tensor of shape (N, C, H, W)     - 'labels': Tensor of shape (N,)     """      def __init__(self, tar_pth_path, transform=None):         self.transform = transform         self.data = []         self.labels = []          # Extract data from .tar.pth file         with open(tar_pth_path, 'rb') as f:             try:                 # Try to load directly                 archive = torch.load(f, map_location='cpu')             except Exception as e:                 print(f"Error loading {tar_pth_path}: {e}")                 raise e          # Assuming the archive is a dict with 'data' and 'labels'         if isinstance(archive, dict):             self.data = archive.get('data', [])             self.labels = archive.get('labels', [])         else:             # If the archive is a list or other structure, adjust accordingly             self.data = archive             self.labels = [0] * len(self.data)  # Dummy labels if not provided          # Convert to tensors if not already         if not isinstance(self.data, torch.Tensor):             self.data = torch.tensor(self.data)         if not isinstance(self.labels, torch.Tensor):             self.labels = torch.tensor(self.labels)      def __len__(self):         return len(self.data)      def __getitem__(self, idx):         image = self.data[idx]         label = self.labels[idx]          if self.transform:             image = self.transform(image)          return image, label  # --------------------- # 3. Feature Extractor # ---------------------  class FeatureExtractor(nn.Module):     """Pre-trained CNN for feature extraction."""      def __init__(self, pretrained=True):         super(FeatureExtractor, self).__init__()         self.model = models.resnet18(pretrained=pretrained)         # Remove the final classification layer         self.features = nn.Sequential(*list(self.model.children())[:-1])      def forward(self, x):         with torch.no_grad():             x = self.features(x)             x = x.view(x.size(0), -1)  # Flatten         return x  # --------------------- # 4. Classifier with Knowledge Distillation # ---------------------  class Classifier(nn.Module):     """Classifier that can incorporate knowledge distillation from a previous model."""      def __init__(self, input_dim, num_classes, previous_model=None, alpha=0.5):         super(Classifier, self).__init__()         self.fc = nn.Linear(input_dim, num_classes)         self.previous_model = previous_model         self.alpha = alpha  # Weight for distillation loss      def forward(self, x):         return self.fc(x)  # --------------------- # 5. Memory Buffer for Experience Replay # ---------------------  class MemoryBuffer:     """Experience Replay Buffer to store representative samples from previous domains."""      def __init__(self, capacity):         self.capacity = capacity         self.buffer = []      def add_samples(self, samples, labels):         for x, y in zip(samples, labels):             if len(self.buffer) >= self.capacity:                 self.buffer.pop(0)  # Remove oldest samples             self.buffer.append((x, y))      def get_samples(self, batch_size):         if len(self.buffer) == 0:             return None, None         selected = random.sample(self.buffer, min(batch_size, len(self.buffer)))         images, labels = zip(*selected)         return torch.stack(images), torch.tensor(labels)  # --------------------- # 6. Random Mixup (RandMix) # ---------------------  def randmix(images, alpha=0.2):     """     Applies mixup augmentation.     Args:         images: Batch of images.         alpha: Mixup parameter.     Returns:         mixed_images: Augmented images.     """     lam = np.random.beta(alpha, alpha)     index = torch.randperm(images.size(0))     mixed_images = lam * images + (1 - lam) * images[index]     return mixed_images, lam, index  # --------------------- # 7. Pseudo-Labeling (T2PL) # ---------------------  def generate_pseudo_labels(model, feature_extractor, dataloader, threshold=0.5):     """     Generates pseudo-labels for unlabeled data using high-confidence predictions.     Args:         model: Current classifier model.         feature_extractor: Feature extractor model.         dataloader: Dataloader for unlabeled data.         threshold: Confidence threshold to select top 50% predictions.     Returns:         pseudo_images: List of images with high-confidence predictions.         pseudo_labels: Corresponding pseudo-labels.     """     model.eval()     feature_extractor.eval()     pseudo_images = []     pseudo_labels = []     confidences = []      all_confidences = []     all_preds = []     all_images = []      with torch.no_grad():         for images, _ in dataloader:             images = images.to(device)             features = feature_extractor(images)             outputs = model(features)             probs = nn.functional.softmax(outputs, dim=1)             max_conf, preds = torch.max(probs, dim=1)             all_confidences.append(max_conf.cpu())             all_preds.append(preds.cpu())             all_images.append(images.cpu())      all_confidences = torch.cat(all_confidences)     all_preds = torch.cat(all_preds)     all_images = torch.cat(all_images)      # Select top 50% confident samples     num_select = all_confidences.size(0) // 2     top_conf, top_indices = torch.topk(all_confidences, num_select)     selected_images = all_images[top_indices]     selected_labels = all_preds[top_indices]      return selected_images, selected_labels, top_conf  # --------------------- # 8. Prototype Contrastive Alignment (PCA) # ---------------------  class PrototypeContrastiveLoss(nn.Module):     """     Prototype Contrastive Loss to align class prototypes across domains.     Encourages similarity between class prototypes and data samples,     while separating different class prototypes.     """      def __init__(self, temperature=0.07):         super(PrototypeContrastiveLoss, self).__init__()         self.temperature = temperature         self.cross_entropy = nn.CrossEntropyLoss()      def forward(self, features, prototypes, labels):         """         Args:             features: Feature representations of samples (batch_size, feature_dim)             prototypes: Class prototypes (num_classes, feature_dim)             labels: Ground truth labels (batch_size)         Returns:             loss: Contrastive loss value         """         # Normalize features and prototypes         features = nn.functional.normalize(features, dim=1)         prototypes = nn.functional.normalize(prototypes, dim=1)          # Compute similarity         logits = torch.matmul(features, prototypes.t()) / self.temperature         loss = self.cross_entropy(logits, labels)         return loss  # --------------------- # 9. Training Function # ---------------------  def train_model(     model,     feature_extractor,     optimizer,     criterion,     dataloader,     memory_buffer,     previous_model=None,     alpha=0.5,     randmix_alpha=0.2,     use_mixup=False ):     """     Trains the classifier model on the given data.     Args:         model: Current classifier model.         feature_extractor: Feature extractor model.         optimizer: Optimizer.         criterion: Loss function (CrossEntropyLoss).         dataloader: Dataloader for current training data.         memory_buffer: Experience replay buffer.         previous_model: Previous classifier model for knowledge distillation.         alpha: Weight for distillation loss.         randmix_alpha: Parameter for mixup.         use_mixup: Whether to apply mixup.     """     model.train()     feature_extractor.eval()  # Freeze feature extractor      for epoch in range(NUM_EPOCHS):         running_loss = 0.0         for images, labels in dataloader:             images = images.to(device)             labels = labels.to(device)              with torch.no_grad():                 features = feature_extractor(images)              if use_mixup:                 # Apply mixup augmentation                 mixed_images, lam, index = randmix(images, alpha=randmix_alpha)                 mixed_features = feature_extractor(mixed_images.to(device))                  outputs = model(mixed_features)                 loss = lam * criterion(outputs, labels.to(device)) + (1 - lam) * criterion(outputs, labels.to(device)[index].to(device))             else:                 outputs = model(features)                 loss = criterion(outputs, labels.to(device))              # Knowledge Distillation Loss             if previous_model is not None:                 with torch.no_grad():                     previous_features = feature_extractor(images)                     previous_outputs = previous_model(previous_features)                 distillation_loss = nn.KLDivLoss()(nn.functional.log_softmax(outputs / 2, dim=1),                                                   nn.functional.softmax(previous_outputs / 2, dim=1))                 loss = (1 - alpha) * loss + alpha * distillation_loss              optimizer.zero_grad()             loss.backward()             optimizer.step()              running_loss += loss.item()          avg_loss = running_loss / len(dataloader)         print(f"Epoch [{epoch+1}/{NUM_EPOCHS}] Loss: {avg_loss:.4f}")      # After training, add samples to memory     # Use a subset of the current batch to add to memory     if use_mixup:         memory_buffer.add_samples(mixed_images.cpu(), labels.cpu())     else:         memory_buffer.add_samples(images.cpu(), labels.cpu())  # --------------------- # 10. Evaluation Function # ---------------------  def evaluate_model(model, feature_extractor, dataloader):     """     Evaluates the model on the given data.     Args:         model: Classifier model.         feature_extractor: Feature extractor model.         dataloader: Dataloader for evaluation data.     Returns:         accuracy: Classification accuracy in percentage.     """     model.eval()     feature_extractor.eval()     all_preds = []     all_labels = []     with torch.no_grad():         for images, labels in dataloader:             images = images.to(device)             labels = labels.to(device)             features = feature_extractor(images)             outputs = model(features)             _, preds = torch.max(outputs, 1)             all_preds.extend(preds.cpu().numpy())             all_labels.extend(labels.cpu().numpy())     accuracy = accuracy_score(all_labels, all_preds) * 100     return accuracy  # --------------------- # 11. Continual Learning Pipeline # ---------------------  def continual_learning_pipeline(     train_tar_pth_paths,     heldout_tar_pth_paths,     num_classes=10,     memory_size=500,     pseudo_label_threshold=0.5,     alpha=0.5,  # Weight for knowledge distillation     randmix_alpha=0.2,     use_mixup=True ):     """     Main pipeline to perform continual learning on sequential domains.     Args:         train_tar_pth_paths: List of paths to training .tar.pth files (D1 to D20).         heldout_tar_pth_paths: List of paths to held-out .tar.pth files (D̂1 to D̂20).         num_classes: Number of classes.         memory_size: Capacity of experience replay buffer.         pseudo_label_threshold: Threshold for selecting high-confidence pseudo-labels.         alpha: Weight for knowledge distillation loss.         randmix_alpha: Parameter for mixup.         use_mixup: Whether to apply mixup.     Returns:         performance_matrix: Numpy array containing accuracies (models x datasets).     """     # Define transformations     transform = transforms.Compose([         transforms.Resize((32, 32)),         transforms.ToTensor(),     ])      # Initialize feature extractor     feature_extractor = FeatureExtractor(pretrained=True).to(device)      # Initialize memory buffer     memory_buffer = MemoryBuffer(memory_size)      # Initialize variables     previous_model = None     performance_matrix = np.zeros((len(train_tar_pth_paths), len(heldout_tar_pth_paths)))      # Iterate through each training dataset     for i, train_path in enumerate(train_tar_pth_paths):         print(f"\n=== Training on Dataset D{i+1} ===")          # Load current training dataset         if i == 0:             # D1 is labeled             current_dataset = TarPTHDataset(train_path, transform=transform)         else:             # D2 to D20 are unlabeled             # We will use pseudo-labeling             current_dataset = TarPTHDataset(train_path, transform=transform)          # Create DataLoader for current dataset         current_loader = DataLoader(current_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)          # Initialize or update classifier         if i == 0:             # Initial classifier             classifier = Classifier(input_dim=512, num_classes=num_classes).to(device)         else:             # Subsequent classifiers with knowledge distillation             classifier = Classifier(input_dim=512, num_classes=num_classes,                                     previous_model=copy.deepcopy(previous_model),                                     alpha=alpha).to(device)          # Define optimizer and loss         optimizer = optim.Adam(classifier.parameters(), lr=LEARNING_RATE)         criterion = nn.CrossEntropyLoss()          if i == 0:             # Task 1: Train on D1 to D10             if i < 10:                 # D1 to D10 share the same distribution                 # Direct supervised training                 train_model(                     model=classifier,                     feature_extractor=feature_extractor,                     optimizer=optimizer,                     criterion=criterion,                     dataloader=current_loader,                     memory_buffer=memory_buffer,                     previous_model=previous_model,                     alpha=alpha,                     randmix_alpha=randmix_alpha,                     use_mixup=use_mixup                 )             else:                 # D11 to D20: different distributions                 # Will handle in Task 2                 pass         else:             if i < 10:                 # Task 1: Training on D2 to D10                 # Apply pseudo-labeling                 pseudo_images, pseudo_labels, confidences = generate_pseudo_labels(                     model=previous_model,                     feature_extractor=feature_extractor,                     dataloader=current_loader,                     threshold=pseudo_label_threshold                 )                  # Create pseudo-labeled dataset                 if len(pseudo_images) > 0:                     pseudo_dataset = torch.utils.data.TensorDataset(pseudo_images, pseudo_labels)                     pseudo_loader = DataLoader(pseudo_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)                     combined_dataset = ConcatDataset([current_dataset, pseudo_dataset])                     combined_loader = DataLoader(combined_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)                 else:                     combined_loader = current_loader                  # Include experience replay from memory                 replay_images, replay_labels = memory_buffer.get_samples(BATCH_SIZE)                 if replay_images is not None:                     replay_dataset = torch.utils.data.TensorDataset(replay_images, replay_labels)                     replay_loader = DataLoader(replay_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)                     final_dataset = ConcatDataset([combined_dataset, replay_dataset])                     final_loader = DataLoader(final_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)                 else:                     final_loader = combined_loader                  # Train the model with combined data                 train_model(                     model=classifier,                     feature_extractor=feature_extractor,                     optimizer=optimizer,                     criterion=criterion,                     dataloader=final_loader,                     memory_buffer=memory_buffer,                     previous_model=previous_model,                     alpha=alpha,                     randmix_alpha=randmix_alpha,                     use_mixup=use_mixup                 )             else:                 # Task 2: Training on D11 to D20                 # Apply domain adaptation techniques                 # Generate pseudo-labels                 pseudo_images, pseudo_labels, confidences = generate_pseudo_labels(                     model=previous_model,                     feature_extractor=feature_extractor,                     dataloader=current_loader,                     threshold=pseudo_label_threshold                 )                  # Create pseudo-labeled dataset                 if len(pseudo_images) > 0:                     pseudo_dataset = torch.utils.data.TensorDataset(pseudo_images, pseudo_labels)                     pseudo_loader = DataLoader(pseudo_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)                     combined_dataset = ConcatDataset([current_dataset, pseudo_dataset])                     combined_loader = DataLoader(combined_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)                 else:                     combined_loader = current_loader                  # Include experience replay from memory                 replay_images, replay_labels = memory_buffer.get_samples(BATCH_SIZE)                 if replay_images is not None:                     replay_dataset = torch.utils.data.TensorDataset(replay_images, replay_labels)                     replay_loader = DataLoader(replay_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)                     final_dataset = ConcatDataset([combined_dataset, replay_dataset])                     final_loader = DataLoader(final_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)                 else:                     final_loader = combined_loader                  # Apply Prototype Contrastive Alignment                 # For simplicity, we'll integrate PCA within the training loop as an additional loss                 # Initialize contrastive loss                 pca_criterion = PrototypeContrastiveLoss()                  # Define a combined training loop                 classifier.train()                 feature_extractor.eval()                  for epoch in range(NUM_EPOCHS):                     running_loss = 0.0                     for images, labels in final_loader:                         images = images.to(device)                         labels = labels.to(device)                          with torch.no_grad():                             features = feature_extractor(images)                          if use_mixup:                             # Apply mixup augmentation                             mixed_images, lam, index = randmix(images, alpha=randmix_alpha)                             mixed_features = feature_extractor(mixed_images.to(device))                              outputs = classifier(mixed_features)                             loss = lam * criterion(outputs, labels.to(device)) + (1 - lam) * criterion(outputs, labels.to(device)[index].to(device))                         else:                             outputs = classifier(features)                             loss = criterion(outputs, labels.to(device))                          # Knowledge Distillation Loss                         if previous_model is not None:                             with torch.no_grad():                                 previous_features = feature_extractor(images)                                 previous_outputs = previous_model(previous_features)                             distillation_loss = nn.KLDivLoss()(nn.functional.log_softmax(outputs / 2, dim=1),                                                               nn.functional.softmax(previous_outputs / 2, dim=1))                             loss = (1 - alpha) * loss + alpha * distillation_loss                          # Prototype Contrastive Loss                         prototypes = compute_prototypes(feature_extractor, classifier, memory_buffer, num_classes)                         contrastive_loss = pca_criterion(features, prototypes, labels)                          # Total loss                         total_loss = loss + contrastive_loss                          optimizer.zero_grad()                         total_loss.backward()                         optimizer.step()                          running_loss += total_loss.item()                      avg_loss = running_loss / len(final_loader)                     print(f"PCA Epoch [{epoch+1}/{NUM_EPOCHS}] Loss: {avg_loss:.4f}")                  # Add current samples to memory                 memory_buffer.add_samples(images.cpu(), labels.cpu())          # Update previous model         previous_model = copy.deepcopy(classifier)          # Evaluate the current model on all relevant held-out datasets         for j in range(i + 1):             heldout_dataset = TarPTHDataset(heldout_tar_pth_paths[j], transform=transform)             heldout_loader = DataLoader(heldout_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)             accuracy = evaluate_model(classifier, feature_extractor, heldout_loader)             performance_matrix[i, j] = accuracy             print(f"Model f{i+1} Accuracy on D̂{j+1}: {accuracy:.2f}%")      return performance_matrix  # --------------------- # 12. Prototype Computation # ---------------------  def compute_prototypes(feature_extractor, classifier, memory_buffer, num_classes=10):     """     Computes class prototypes from the memory buffer.     Args:         feature_extractor: Feature extractor model.         classifier: Current classifier model.         memory_buffer: Experience replay buffer.         num_classes: Number of classes.     Returns:         prototypes: Tensor of shape (num_classes, feature_dim)     """     feature_extractor.eval()     classifier.eval()     class_features = {i: [] for i in range(num_classes)}      with torch.no_grad():         for img, label in memory_buffer.buffer:             img = img.to(device).unsqueeze(0)             label = label.item()             feat = feature_extractor(img)             class_features[label].append(feat.cpu())      prototypes = []     for i in range(num_classes):         if len(class_features[i]) > 0:             # Compute mean feature for the class             class_feat = torch.cat(class_features[i], dim=0)             mean_feat = class_feat.mean(dim=0)             prototypes.append(mean_feat)         else:             # If no samples for the class, use zeros             prototypes.append(torch.zeros(512))      prototypes = torch.stack(prototypes).to(device)     return prototypes  # --------------------- # 13. Preparing Datasets # ---------------------  def prepare_datasets(train_folder, heldout_folder, num_datasets=20):     """     Prepares lists of paths to training and held-out datasets.     Args:         train_folder: Directory containing training .tar.pth files (D1 to D20).         heldout_folder: Directory containing held-out .tar.pth files (D̂1 to D̂20).         num_datasets: Number of datasets (default 20).     Returns:         train_tar_pth_paths: List of training dataset paths.         heldout_tar_pth_paths: List of held-out dataset paths.     """     train_tar_pth_paths = []     heldout_tar_pth_paths = []      for i in range(1, num_datasets + 1):         train_path = os.path.join(train_folder,f"train_data/{i}_train_data.tar.pth")         heldout_path = os.path.join(heldout_folder, f"eval_data/{i}_eval_data.tar.pth")  # Using 'Dhat' to denote held-out         if not os.path.exists(train_path):             raise FileNotFoundError(f"Training dataset {train_path} not found.")         if not os.path.exists(heldout_path):             raise FileNotFoundError(f"Held-out dataset {heldout_path} not found.")         train_tar_pth_paths.append(train_path)         heldout_tar_pth_paths.append(heldout_path)      return train_tar_pth_paths, heldout_tar_pth_paths  # --------------------- # 14. Hyperparameters # ---------------------  BATCH_SIZE = 64 LEARNING_RATE = 1e-3 NUM_EPOCHS = 10 MEMORY_SIZE = 500 PSEUDO_LABEL_THRESHOLD = 0.5 ALPHA = 0.5  # Weight for knowledge distillation RANDMIX_ALPHA = 0.2 USE_MIXUP = True  # --------------------- # 15. Main Execution # ---------------------  if __name__ == "__main__":      #   # Paths to zip files     # train_zip_path = "/content/train_data.zip"     # heldout_zip_path = "/content/eval_data.zip"      # # # Directory where you want to extract the files     # # extracted_train_folder = "path_to_extracted_train_folder"     # # extracted_heldout_folder = "path_to_extracted_heldout_folder"      # extracted_train_folder = "./train_data"      # extracted_heldout_folder = "./heldout_data"       # # Extract train folder     # if not os.path.exists(extracted_train_folder):     #     with zipfile.ZipFile(train_zip_path, 'r') as zip_ref:     #         zip_ref.extractall(extracted_train_folder)      # # Extract heldout folder     # if not os.path.exists(extracted_heldout_folder):     #     with zipfile.ZipFile(heldout_zip_path, 'r') as zip_ref:     #         zip_ref.extractall(extracted_heldout_folder)     # Paths to the directories containing training and held-out .tar.pth files     # Replace 'path_to_train_folder' and 'path_to_heldout_folder' with actual paths     train_folder ="/content/train_data"  # Example: "./datasets/train/"     heldout_folder = "/content/heldout_data"  # Example: "./datasets/heldout/"      # Prepare dataset paths     try:         train_tar_pth_paths, heldout_tar_pth_paths = prepare_datasets(train_folder, heldout_folder)     except FileNotFoundError as e:         print(e)         exit(1)      # Run continual learning pipeline     performance = continual_learning_pipeline(         train_tar_pth_paths=train_tar_pth_paths,         heldout_tar_pth_paths=heldout_tar_pth_paths,         num_classes=10,         memory_size=MEMORY_SIZE,         pseudo_label_threshold=PSEUDO_LABEL_THRESHOLD,         alpha=ALPHA,         randmix_alpha=RANDMIX_ALPHA,         use_mixup=USE_MIXUP     )      # Display performance matrix     df = pd.DataFrame(         performance,         index=[f"f{i+1}" for i in range(len(train_tar_pth_paths))],         columns=[f"Dhat{j+1}" for j in range(len(heldout_tar_pth_paths))]     )     print("\n=== Performance Matrix (Accuracy %) ===")     print(df)      # Save performance matrix to CSV     df.to_csv("performance_matrix.csv")     print("\nPerformance matrix saved to 'performance_matrix.csv'"), here the training dataset 1 is labelled, others are unlabelled , now trian the model using dataset 1, then use it to predict labels for the dataset  2, and so on , so modify this code a little to accomodate that