Create train.py

train.py

@@ -0,0 +1,263 @@
+# remember to run preprocess.py before training
+# preprocess while training is not as effecient
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import MultiheadAttention
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader, random_split
+import json
+import time
+import os
+import h5py
+import numpy as np
+from tqdm import tqdm
+
+class AttentionBlock(nn.Module):
+    def __init__(self, input_dim, num_heads, key_dim, ff_dim, rate=0.1):
+        super(AttentionBlock, self).__init__()
+        self.multihead_attn = MultiheadAttention(embed_dim=input_dim, num_heads=num_heads)
+        self.dropout1 = nn.Dropout(rate)
+        self.layer_norm1 = nn.LayerNorm(input_dim, eps=1e-6)
+
+        self.ffn = nn.Sequential(
+            nn.Linear(input_dim, ff_dim),
+            nn.ReLU(),
+            nn.Dropout(rate),
+            nn.Linear(ff_dim, input_dim),
+            nn.Dropout(rate)
+        )
+        self.layer_norm2 = nn.LayerNorm(input_dim, eps=1e-6)
+
+    def forward(self, x):
+        attn_output, _ = self.multihead_attn(x, x, x)
+        attn_output = self.dropout1(attn_output)
+        out1 = self.layer_norm1(x + attn_output)
+
+        ffn_output = self.ffn(out1)
+        out2 = self.layer_norm2(out1 + ffn_output)
+
+        return out2
+
+class TextureContrastClassifier(nn.Module):
+    def __init__(self, input_shape, num_heads=4, key_dim=64, ff_dim=256, rate=0.5):
+        super(TextureContrastClassifier, self).__init__()
+        input_dim = input_shape[1]  # assuming the input shape is (seq_len, feature_dim)
+
+        self.rich_texture_attention = AttentionBlock(input_dim, num_heads, key_dim, ff_dim, rate)
+        self.poor_texture_attention = AttentionBlock(input_dim, num_heads, key_dim, ff_dim, rate)
+
+        self.rich_texture_dense = nn.Sequential(
+            nn.Linear(input_dim, 128),
+            nn.ReLU(),
+            nn.Dropout(rate)
+        )
+
+        self.poor_texture_dense = nn.Sequential(
+            nn.Linear(input_dim, 128),
+            nn.ReLU(),
+            nn.Dropout(rate)
+        )
+
+        self.fc = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(input_shape[0] * 128, 256),
+            nn.ReLU(),
+            nn.Dropout(rate),
+            nn.Linear(256, 128),
+            nn.ReLU(),
+            nn.Dropout(rate),
+            nn.Linear(128, 64),
+            nn.ReLU(),
+            nn.Dropout(rate),
+            nn.Linear(64, 32),
+            nn.ReLU(),
+            nn.Dropout(rate),
+            nn.Linear(32, 16),
+            nn.ReLU(),
+            nn.Dropout(rate),
+            nn.Linear(16, 1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, rich_texture, poor_texture):
+        rich_texture = self.rich_texture_attention(rich_texture)
+        rich_texture = self.rich_texture_dense(rich_texture)
+
+        poor_texture = self.poor_texture_attention(poor_texture)
+        poor_texture = self.poor_texture_dense(poor_texture)
+
+        difference = rich_texture - poor_texture
+        output = self.fc(difference)
+
+        return output
+
+import os
+import h5py
+import numpy as np
+from tqdm import tqdm
+
+def load_and_split_data(h5_dir, train_ratio=0.8,max_num=40):
+    train_rich, train_poor, train_labels = [], [], []
+    test_rich, test_poor, test_labels = [], [], []
+
+    for file_name in tqdm(os.listdir(h5_dir)[:60]):
+        if file_name.endswith('.h5'):
+            file_path = os.path.join(h5_dir, file_name)
+            try:
+                with h5py.File(file_path, 'r') as h5f:
+                    rich = h5f['rich'][:]
+                    poor = h5f['poor'][:]
+                    labels = h5f['labels'][:]
+
+                dataset_size = len(labels)
+                train_size = int(train_ratio * dataset_size)
+                indices = np.random.permutation(dataset_size)
+                train_indices = indices[:train_size]
+                test_indices = indices[train_size:]
+
+                train_rich.append(rich[train_indices])
+                train_poor.append(poor[train_indices])
+                train_labels.append(labels[train_indices])
+
+                test_rich.append(rich[test_indices])
+                test_poor.append(poor[test_indices])
+                test_labels.append(labels[test_indices])
+
+            except Exception as e:
+                print(f"Error processing {file_name}: {e}")
+
+    train_rich = np.concatenate(train_rich, axis=0)
+    train_poor = np.concatenate(train_poor, axis=0)
+    train_labels = np.concatenate(train_labels, axis=0)
+
+    test_rich = np.concatenate(test_rich, axis=0)
+    test_poor = np.concatenate(test_poor, axis=0)
+    test_labels = np.concatenate(test_labels, axis=0)
+
+    return train_rich, train_poor, train_labels, test_rich, test_poor, test_labels
+
+class TextureDataset(Dataset):
+    def __init__(self, rich, poor, labels):
+        self.rich = rich
+        self.poor = poor
+        self.labels = labels
+
+    def __len__(self):
+        return len(self.labels)
+
+    def __getitem__(self, idx):
+        rich = torch.tensor(self.rich[idx], dtype=torch.float32)
+        poor = torch.tensor(self.poor[idx], dtype=torch.float32)
+        label = torch.tensor(self.labels[idx], dtype=torch.float32)
+        return rich, poor, label
+
+def validate(model, test_loader, criterion, device):
+    model.eval()
+    val_loss = 0.0
+    correct = 0
+    total = 0
+
+    with torch.no_grad():
+        for rich, poor, labels in test_loader:
+            rich, poor, labels = rich.to(device), poor.to(device), labels.to(device)
+
+            outputs = model(rich, poor)
+            outputs = outputs.squeeze()
+
+            loss = criterion(outputs, labels)
+            val_loss += loss.item()
+
+            predicted = (outputs > 0.5).float()
+            total += labels.size(0)
+            correct += (predicted == labels).sum().item()
+
+    val_loss /= len(test_loader)
+    val_accuracy = correct / total
+    print(f'Validation Loss: {val_loss:.4f}, Validation Accuracy: {val_accuracy:.4f}')
+    return val_loss, val_accuracy
+
+
+
+h5_dir = '/content/drive/MyDrive/h5saves'
+train_rich, train_poor, train_labels, test_rich, test_poor, test_labels = load_and_split_data(h5_dir, train_ratio=0.8)
+print(f"Training data: {len(train_labels)} samples")
+print(f"Testing data: {len(test_labels)} samples")
+train_dataset = TextureDataset(train_rich, train_poor, train_labels)
+test_dataset = TextureDataset(test_rich, test_poor, test_labels)
+batch_size = 2048
+train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
+test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)
+
+input_shape = (128, 256)
+model = TextureContrastClassifier(input_shape)
+criterion = nn.BCELoss()
+optimizer = optim.Adam(model.parameters(), lr=0.0001)
+scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=5, verbose=True)
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model.to(device)
+
+history = {'train_loss': [], 'val_loss': [], 'train_accuracy':[], 'val_accuracy': []}
+save_dir = '/content/drive/MyDrive/model_checkpoints'
+if not os.path.exists(save_dir):
+    os.makedirs(save_dir)
+num_epochs = 100
+
+
+
+for epoch in range(num_epochs):
+    model.train()
+    running_loss = 0.0
+    correct = 0
+    total = 0
+
+    batch_loss = 0.0
+
+    for batch_idx, (rich, poor, labels) in enumerate(train_loader):
+        rich, poor, labels = rich.to(device), poor.to(device), labels.to(device)
+
+        optimizer.zero_grad()
+
+        outputs = model(rich, poor)
+        outputs = outputs.squeeze()
+
+        loss = criterion(outputs, labels)
+        loss.backward()
+        optimizer.step()
+
+        running_loss += loss.item()
+        batch_loss += loss.item()
+
+        predicted = (outputs > 0.5).float()
+        total += labels.size(0)
+        correct += (predicted == labels).sum().item()
+
+        if (batch_idx + 1) % 5 == 0:
+            print(f'\rEpoch [{epoch+1}/{num_epochs}], Batch [{batch_idx+1}], Loss: {batch_loss / 5:.4f}, Accuracy: {correct / total:.2f}', end='')
+            batch_loss = 0.0
+
+    avg_train_loss = running_loss / len(train_loader)
+    train_accuracy =  correct / total
+
+    val_loss, val_accuracy = validate(model, test_loader, criterion, device)
+
+    history['train_loss'].append(avg_train_loss)
+    history['val_loss'].append(val_loss)
+    history['val_accuracy'].append(val_accuracy)
+    history['train_accuracy'].append(train_accuracy)
+
+    scheduler.step(val_loss)
+
+    checkpoint_path = os.path.join(save_dir, f'model_epoch_{epoch+1}.pth')
+    torch.save(model.state_dict(), checkpoint_path)
+    print(f'\nModel checkpoint saved for epoch {epoch+1}')
+
+    print(f'Epoch [{epoch+1}/{num_epochs:.4f}], Training Loss: {avg_train_loss:.4f}, Training Accuracy: {train_accuracy:.4f} Validation Loss: {val_loss:.4f}, Validation Accuracy: {val_accuracy:.4f}')
+
+history_path = os.path.join(save_dir, 'training_history.json')
+with open(history_path, 'w') as f:
+    json.dump(history, f)
+
+print('Finished Training')
+print(f'Training history saved at {history_path}')