Update utils/pretraining.py

utils/pretraining.py

@@ -1,150 +1,273 @@
-#%% PACKAGES & MODULES
-import torch
-import torch.nn as nn
-import torch.optim as optim
-from torch.optim.lr_scheduler import StepLR
-from inference import prepare_for_lwm
-from input_preprocess import tokenizer
-from lwm_model import lwm
-import numpy as np
-
-#%% PARAMETERS
-n_epochs = 100
-n_layers = 12
-n_heads = 12
-d_model = 64
-d_ff = d_model * 4
-d_k = d_model // n_heads
-d_v = d_model // n_heads
-dropout = 0.1
-max_len = 129
-element_length = 16
-batch_size = 64
-train_ratio = 0.7
-val_ratio = 0.2
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
-
-#%% PRE-TRAINING DATA GENERATION
-# The following DeepMIMO scenarios are not enough for pre-training a 
-# Transformer-based foundation model like LWM. Add more scenarios for 
-# more effective pre-training. The instruction for reproducing the actual 
-# dataset used for pre-training LWM can be found in the Huggingface forum.
-scenario_names = np.array([
-    "city_18_denver", "city_15_indianapolis", "city_19_oklahoma", 
-    "city_12_fortworth", "city_11_santaclara", "city_7_sandiego"
-])
-
-scenario_idxs = np.array([0, 1, 2, 3, 4, 5])  
-selected_scenario_names = scenario_names[scenario_idxs]
-
-preprocessed_chs = tokenizer(
-    selected_scenario_names=selected_scenario_names, 
-    manual_data=None, 
-    gen_raw=False) 
-
-#%% DATALOADER
-train_size = int(train_ratio * len(preprocessed_chs))
-val_size = int(val_ratio * len(preprocessed_chs))
-test_size = len(preprocessed_chs) - val_size - train_size
-
-train_data, val_data, test_data = torch.utils.data.random_split(
-    preprocessed_chs, [train_size, val_size, test_size]
-)
-
-train_loader = prepare_for_lwm(train_data, device, batch_size=batch_size, shuffle=True)
-val_loader = prepare_for_lwm(val_data, device, batch_size=batch_size, shuffle=True)
-test_loader = prepare_for_lwm(test_data, device, batch_size=batch_size, shuffle=True)
-
-# %% Model
-load_model = False
-
-model = lwm()
-model.to(device)
-
-if load_model:
-    model_name = 'models/pretrained_model.pth'
-    model.load_state_dict(torch.load(model_name))
-    print(f"Model loaded from {model_name}")
-    
-# Loss function
-criterionMLM = nn.MSELoss()
-
-# %% Optimizer and Scheduler
-adaptive_lr = False
-
-optimizer = optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-5)
-scheduler = (
-    optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min')
-    if adaptive_lr
-    else StepLR(optimizer, step_size=10, gamma=0.9)
-)
-
-# %% Training
-training_loss = []
-validation_loss = []
-
-def train(model, dataloader, optimizer, scheduler=None, device="cuda"):
-
-    model.train()
-    running_loss = 0.0
-    criterionMCM = nn.MSELoss()
-
-    for idx, batch in enumerate(dataloader):
-        input_ids = batch[0].to(device)
-        masked_tokens = batch[1].to(device)
-        masked_pos = batch[2].to(device)
-        
-        optimizer.zero_grad()
-        
-        logits_lm, _ = model(input_ids, masked_pos)
-        loss_lm = criterionMCM(logits_lm, masked_tokens)
-        loss = loss_lm / torch.var(masked_tokens) 
-        
-        loss.backward()
-        optimizer.step()
-
-        if scheduler is not None:
-            scheduler.step()
-
-        running_loss += loss.item()
-
-    average_loss = running_loss / len(dataloader)
-
-    return average_loss
-
-def validate(model, dataloader, device="cuda"):
-    model.eval()
-    running_loss = 0.0
-    criterionMCM = nn.MSELoss()
-
-    with torch.no_grad():
-        for idx, batch in enumerate(dataloader):
-            input_ids = batch[0].to(device)
-            masked_tokens = batch[1].to(device)
-            masked_pos = batch[2].to(device)
-
-            logits_lm, _ = model(input_ids, masked_pos)
-
-            loss_lm = criterionMCM(logits_lm, masked_tokens)
-            loss = loss_lm / torch.var(masked_tokens)  
-
-            running_loss += loss.item()
-
-    average_loss = running_loss / len(dataloader)
-
-    return average_loss
-
-# %% Training Loop
-for epoch in range(n_epochs):
-    print(f"Epoch {epoch + 1}/{n_epochs}")
-
-    # Training step
-    train_loss = train(model, train_loader, optimizer, scheduler, device)
-    training_loss.append(train_loss)
-    print(f"Training Loss: {train_loss:.4f}")
-
-    # Validation step
-    if val_loader is not None:
-        val_loss = validate(model, val_loader, device)
-        validation_loss.append(val_loss)
-        print(f"Validation Loss: {val_loss:.4f}")
+#%% PACKAGES & MODULES
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.optim.lr_scheduler import StepLR
+from inference import prepare_for_lwm
+from input_preprocess import tokenizer
+from lwm_model import lwm
+import numpy as np
+import DeepMIMOv3
+
+#%% PRE-TRAINING SCENARIO CONFIG
+def get_parameters(scenario):
+    
+    n_ant_bs = 32
+    n_ant_ue = 1
+    n_subcarriers = 32 
+    scs = 30e3
+    
+    row_column_users = {
+    'asu_campus1': {
+        'n_rows': 321,
+        'n_per_row': 411
+    },
+    'Boston5G_3p5': {
+        'n_rows': [812,1622],
+        'n_per_row': 595 
+    },
+    'city_0_newyork': {
+        'n_rows': 44,
+        'n_per_row': 117
+    },
+    'city_1_losangeles': {
+        'n_rows': 57,
+        'n_per_row': 81
+    },
+    'city_2_chicago': {
+        'n_rows': 56,
+        'n_per_row': 80
+    },
+    'city_3_houston': {
+        'n_rows': 62,
+        'n_per_row': 81
+    },
+    'city_4_phoenix': {
+        'n_rows': 79,
+        'n_per_row': 86
+    },
+    'city_5_philadelphia': {
+        'n_rows': 96,
+        'n_per_row': 66
+    },
+    'city_6_miami': {
+        'n_rows': 80,
+        'n_per_row': 87 
+    },
+    'city_8_dallas': {
+        'n_rows': 83,
+        'n_per_row': 76
+    },
+    'city_9_sanfrancisco': {
+        'n_rows': 79,
+        'n_per_row': 83
+    },
+    'city_10_austin': {
+        'n_rows': 102,
+        'n_per_row': 55
+    },
+    'city_13_columbus': {
+        'n_rows': 71,
+        'n_per_row': 96
+    },
+    'city_17_seattle': {
+        'n_rows': 74,
+        'n_per_row': 82
+    },
+    'O1_3p5': {
+        'n_rows': 5203,
+        'n_per_row': 181
+    },
+    'city_18_denver': {
+        'n_rows': 85,
+        'n_per_row': 82
+    },
+    'city_15_indianapolis': {
+        'n_rows': 80,
+        'n_per_row': 79
+    },
+    'city_19_oklahoma': {
+        'n_rows': 82,
+        'n_per_row': 75
+    },
+    'city_12_fortworth': {
+        'n_rows': 86,
+        'n_per_row': 72
+    },
+    'city_11_santaclara': {
+        'n_rows': 47,
+        'n_per_row': 114
+    },
+    'city_7_sandiego': {
+        'n_rows': 71,
+        'n_per_row': 83
+    }}
+    
+    parameters = DeepMIMOv3.default_params()
+    parameters['dataset_folder'] = './scenarios'
+    parameters['scenario'] = scenario
+    
+    if scenario == 'O1_3p5':
+        parameters['active_BS'] = np.array([4])
+    elif scenario in ['city_14_charlotte', 'city_18_denver', 'city_15_indianapolis']:
+        parameters['active_BS'] = np.array([3])
+    else:
+        parameters['active_BS'] = np.array([1])
+        
+    if scenario == 'Boston5G_3p5':
+        parameters['user_rows'] = np.arange(row_column_users[scenario]['n_rows'][0],
+                                            row_column_users[scenario]['n_rows'][1])
+    else:
+        parameters['user_rows'] = np.arange(row_column_users[scenario]['n_rows'])
+    parameters['bs_antenna']['shape'] = np.array([n_ant_bs, 1]) # Horizontal, Vertical 
+    parameters['bs_antenna']['rotation'] = np.array([0,0,-135]) # (x,y,z)
+    parameters['ue_antenna']['shape'] = np.array([n_ant_ue, 1])
+    parameters['enable_BS2BS'] = False
+    parameters['OFDM']['subcarriers'] = n_subcarriers
+    parameters['OFDM']['selected_subcarriers'] = np.arange(n_subcarriers)
+    
+    parameters['OFDM']['bandwidth'] = scs * n_subcarriers / 1e9
+    parameters['num_paths'] = 20
+    
+    return parameters, row_column_users, n_ant_bs, n_ant_ue, n_subcarriers
+    
+#%% PARAMETERS
+n_epochs = 100
+n_layers = 12
+n_heads = 12
+d_model = 64
+d_ff = d_model * 4
+d_k = d_model // n_heads
+d_v = d_model // n_heads
+dropout = 0.1
+max_len = 129
+element_length = 16
+batch_size = 64
+train_ratio = 0.7
+val_ratio = 0.2
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+#%% PRE-TRAINING DATA GENERATION
+# The following DeepMIMO scenarios are not enough for pre-training a 
+# Transformer-based foundation model like LWM. Add more scenarios for 
+# more effective pre-training. The instruction for reproducing the actual 
+# dataset used for pre-training LWM can be found in the Huggingface forum.
+scenario_names = np.array([
+    "city_18_denver", "city_15_indianapolis", "city_19_oklahoma", 
+    "city_12_fortworth", "city_11_santaclara", "city_7_sandiego"
+])
+
+scenario_idxs = np.array([0, 1, 2, 3, 4, 5])  
+selected_scenario_names = scenario_names[scenario_idxs]
+
+preprocessed_chs = tokenizer(
+    selected_scenario_names=selected_scenario_names, 
+    manual_data=None, 
+    gen_raw=False) 
+
+#%% DATALOADER
+train_size = int(train_ratio * len(preprocessed_chs))
+val_size = int(val_ratio * len(preprocessed_chs))
+test_size = len(preprocessed_chs) - val_size - train_size
+
+train_data, val_data, test_data = torch.utils.data.random_split(
+    preprocessed_chs, [train_size, val_size, test_size]
+)
+
+train_loader = prepare_for_lwm(train_data, device, batch_size=batch_size, shuffle=True)
+val_loader = prepare_for_lwm(val_data, device, batch_size=batch_size, shuffle=True)
+test_loader = prepare_for_lwm(test_data, device, batch_size=batch_size, shuffle=True)
+
+# %% Model
+load_model = False
+
+model = lwm()
+model.to(device)
+
+if load_model:
+    model_name = 'models/pretrained_model.pth'
+    model.load_state_dict(torch.load(model_name))
+    print(f"Model loaded from {model_name}")
+    
+# Loss function
+criterionMLM = nn.MSELoss()
+
+# %% Optimizer and Scheduler
+adaptive_lr = False
+
+optimizer = optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-5)
+scheduler = (
+    optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min')
+    if adaptive_lr
+    else StepLR(optimizer, step_size=10, gamma=0.9)
+)
+
+# %% Training
+training_loss = []
+validation_loss = []
+
+def train(model, dataloader, optimizer, scheduler=None, device="cuda"):
+
+    model.train()
+    running_loss = 0.0
+    criterionMCM = nn.MSELoss()
+
+    for idx, batch in enumerate(dataloader):
+        input_ids = batch[0].to(device)
+        masked_tokens = batch[1].to(device)
+        masked_pos = batch[2].to(device)
+        
+        optimizer.zero_grad()
+        
+        logits_lm, _ = model(input_ids, masked_pos)
+        loss_lm = criterionMCM(logits_lm, masked_tokens)
+        loss = loss_lm / torch.var(masked_tokens) 
+        
+        loss.backward()
+        optimizer.step()
+
+        if scheduler is not None:
+            scheduler.step()
+
+        running_loss += loss.item()
+
+    average_loss = running_loss / len(dataloader)
+
+    return average_loss
+
+def validate(model, dataloader, device="cuda"):
+    model.eval()
+    running_loss = 0.0
+    criterionMCM = nn.MSELoss()
+
+    with torch.no_grad():
+        for idx, batch in enumerate(dataloader):
+            input_ids = batch[0].to(device)
+            masked_tokens = batch[1].to(device)
+            masked_pos = batch[2].to(device)
+
+            logits_lm, _ = model(input_ids, masked_pos)
+
+            loss_lm = criterionMCM(logits_lm, masked_tokens)
+            loss = loss_lm / torch.var(masked_tokens)  
+
+            running_loss += loss.item()
+
+    average_loss = running_loss / len(dataloader)
+
+    return average_loss
+
+# %% Training Loop
+for epoch in range(n_epochs):
+    print(f"Epoch {epoch + 1}/{n_epochs}")
+
+    # Training step
+    train_loss = train(model, train_loader, optimizer, scheduler, device)
+    training_loss.append(train_loss)
+    print(f"Training Loss: {train_loss:.4f}")
+
+    # Validation step
+    if val_loader is not None:
+        val_loss = validate(model, val_loader, device)
+        validation_loss.append(val_loss)
+        print(f"Validation Loss: {val_loss:.4f}")