Update model.py

model.py

@@ -6,12 +6,13 @@ import torch.nn.functional as F
 from typing import List, Dict, Any, Optional
 from collections import defaultdict
 from accelerate import Accelerator
+from torch.utils.data import DataLoader, TensorDataset
+from torch.cuda.amp import GradScaler, autocast
 
 class DynamicModel(nn.Module):
     def __init__(self, sections: Dict[str, List[Dict[str, Any]]]):
         super(DynamicModel, self).__init__()
         self.sections = nn.ModuleDict()
-
         if not sections:
             sections = {
                 'default': [{
@@ -22,7 +23,6 @@ class DynamicModel(nn.Module):
                     'dropout': 0.1
                 }]
             }
-
         for section_name, layers in sections.items():
             self.sections[section_name] = nn.ModuleList()
             for layer_params in layers:
@@ -32,10 +32,8 @@ class DynamicModel(nn.Module):
     def create_layer(self, layer_params: Dict[str, Any]) -> nn.Module:
         layers = []
         layers.append(nn.Linear(layer_params['input_size'], layer_params['output_size']))
-
         if layer_params.get('batch_norm', False):
             layers.append(nn.BatchNorm1d(layer_params['output_size']))
-
         activation = layer_params.get('activation', 'relu')
         if activation == 'relu':
             layers.append(nn.ReLU(inplace=True))
@@ -49,23 +47,17 @@ class DynamicModel(nn.Module):
             layers.append(nn.ELU(alpha=1.0, inplace=True))
         elif activation is not None:
             raise ValueError(f"Unsupported activation function: {activation}")
-
         if dropout_rate := layer_params.get('dropout', 0.0):
             layers.append(nn.Dropout(p=dropout_rate))
-
         if hidden_layers := layer_params.get('hidden_layers', []):
             for hidden_layer_params in hidden_layers:
                 layers.append(self.create_layer(hidden_layer_params))
-
         if layer_params.get('memory_augmentation', True):
             layers.append(MemoryAugmentationLayer(layer_params['output_size']))
-
         if layer_params.get('hybrid_attention', True):
             layers.append(HybridAttentionLayer(layer_params['output_size']))
-
         if layer_params.get('dynamic_flash_attention', True):
             layers.append(DynamicFlashAttentionLayer(layer_params['output_size']))
-
         return nn.Sequential(*layers)
 
     def forward(self, x: torch.Tensor, section_name: Optional[str] = None) -> torch.Tensor:
@@ -111,38 +103,30 @@ class DynamicFlashAttentionLayer(nn.Module):
 def parse_xml_file(file_path: str) -> List[Dict[str, Any]]:
     tree = ET.parse(file_path)
     root = tree.getroot()
-
     layers = []
     for layer in root.findall('.//layer'):
         layer_params = {}
         layer_params['input_size'] = int(layer.get('input_size', 128))
         layer_params['output_size'] = int(layer.get('output_size', 256))
         layer_params['activation'] = layer.get('activation', 'relu').lower()
-
         if layer_params['activation'] not in ['relu', 'tanh', 'sigmoid', 'none']:
             raise ValueError(f"Unsupported activation function: {layer_params['activation']}")
-
         if layer_params['input_size'] <= 0 or layer_params['output_size'] <= 0:
             raise ValueError("Layer dimensions must be positive integers")
-
         layers.append(layer_params)
-
     if not layers:
         layers.append({
             'input_size': 128,
             'output_size': 256,
             'activation': 'relu'
         })
-
     return layers
 
 def create_model_from_folder(folder_path: str) -> DynamicModel:
     sections = defaultdict(list)
-
     if not os.path.exists(folder_path):
         print(f"Warning: Folder {folder_path} does not exist. Creating model with default configuration.")
         return DynamicModel({})
-
     xml_files_found = False
     for root, dirs, files in os.walk(folder_path):
         for file in files:
@@ -155,64 +139,59 @@ def create_model_from_folder(folder_path: str) -> DynamicModel:
                     sections[section_name].extend(layers)
                 except Exception as e:
                     print(f"Error processing {file_path}: {str(e)}")
-
     if not xml_files_found:
         print("Warning: No XML files found. Creating model with default configuration.")
         return DynamicModel({})
-
     return DynamicModel(dict(sections))
 
 def main():
     folder_path = 'data'
     model = create_model_from_folder(folder_path)
-
     print(f"Created dynamic PyTorch model with sections: {list(model.sections.keys())}")
-
     # Print the model architecture
     print(model)
-
     first_section = next(iter(model.sections.keys()))
     first_layer = model.sections[first_section][0]
     input_features = first_layer[0].in_features
-
     # Ensure the input tensor size matches the expected input size
     sample_input = torch.randn(1, input_features)
     output = model(sample_input)
     print(f"Sample output shape: {output.shape}")
 
     accelerator = Accelerator()
-
     optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
     criterion = nn.CrossEntropyLoss()
     num_epochs = 10
-
-    dataset = torch.utils.data.TensorDataset(
+    dataset = TensorDataset(
         torch.randn(100, input_features),
         torch.randint(0, 2, (100,))
     )
-    train_dataloader = torch.utils.data.DataLoader(
-        dataset, 
-        batch_size=16, 
+    train_dataloader = DataLoader(
+        dataset,
+        batch_size=8,  # Reduced batch size
         shuffle=True
     )
 
     model, optimizer, train_dataloader = accelerator.prepare(
-        model, 
-        optimizer, 
+        model,
+        optimizer,
         train_dataloader
     )
 
+    scaler = GradScaler()  # Mixed precision training
+
     for epoch in range(num_epochs):
         model.train()
         total_loss = 0
         for batch_idx, (inputs, labels) in enumerate(train_dataloader):
             optimizer.zero_grad()
-            outputs = model(inputs)
-            loss = criterion(outputs, labels)
-            accelerator.backward(loss)
-            optimizer.step()
+            with autocast():  # Mixed precision training
+                outputs = model(inputs)
+                loss = criterion(outputs, labels)
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
             total_loss += loss.item()
-
         avg_loss = total_loss / len(train_dataloader)
         print(f"Epoch {epoch+1}/{num_epochs}, Average Loss: {avg_loss:.4f}")