upload side scripts

utils/beamforming.py

@@ -0,0 +1,67 @@
+#%% PACKAGES & MODEULS
+import numpy as np
+import torch
+from input_preprocess import DeepMIMO_data_gen, deepmimo_data_cleaning, tokenizer
+from inference import lwm_inference, create_raw_dataset
+from lwm_model import lwm
+
+#%% DEEPMIMO DATA GENERATION
+scenario_names = np.array([
+    "city_18_denver", "city_15_indianapolis", "city_19_oklahoma", 
+    "city_12_fortworth", "city_11_santaclara", "city_7_sandiego"
+])
+
+bf_scenario_idx = 3
+scenario_idxs = np.array([bf_scenario_idx])  
+selected_scenario_names = scenario_names[scenario_idxs]
+
+deepmimo_data = [DeepMIMO_data_gen(scenario_name) for scenario_name in selected_scenario_names]
+cleaned_deepmimo_data = [deepmimo_data_cleaning(deepmimo_data[scenario_idx]) for scenario_idx in range(len(deepmimo_data))]
+
+#%% FUNCTION FOR MRT BEAMFORMING
+def compute_mrt_beamforming(channel_data, snr_db=None):
+
+    channel_data = torch.tensor(channel_data[0])
+    mrt_vectors = []
+    snr_linear = 10 ** (snr_db / 10) if snr_db is not None else None
+
+    for idx in range(channel_data.shape[0]):
+        channel = channel_data[idx, 0, :, :]  # Shape: (32, 32)
+
+        if snr_db is not None:
+            # Add complex Gaussian noise to the channel
+            noise_power = torch.mean(torch.abs(channel) ** 2) / snr_linear
+            noise = torch.sqrt(noise_power / 2) * (
+                torch.randn_like(channel) + 1j * torch.randn_like(channel)
+            )
+            channel = channel + noise
+
+        # Compute MRT beamforming vector for each user
+        h_avg = torch.mean(channel, dim=1, keepdim=True)  # Shape: (32, 1)
+        h_conj = torch.conj(h_avg)  # Conjugate of averaged channel vector
+        mrt_vector = h_conj / torch.norm(h_conj, dim=0, keepdim=True)  # Normalize
+
+        mrt_vectors.append(mrt_vector)
+
+    return torch.stack(mrt_vectors, dim=0)  # Shape: (N, 32, 1)
+
+#%% GENERATE BEAMFORMING VECTORS
+beamforming_vectors = compute_mrt_beamforming(cleaned_deepmimo_data)
+
+#%% GENERATE LWM EMBEDDINGS FROM MASKED INPUT CHANNELS
+preprocessed_chs = tokenizer(
+    selected_scenario_names=selected_scenario_names, 
+    manual_data=None, 
+    gen_raw=False) # gen_raw=False masks 15% of the input patches, and LWM will act as a denoiser
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+print(f"Loading the LWM model on {device} ...")
+model = lwm.from_pretrained(device=device)
+
+input_types = ['cls_emb', 'channel_emb', 'raw']
+selected_input_type = input_types[1] 
+
+if selected_input_type in ['cls_emb', 'channel_emb']:
+    dataset = lwm_inference(preprocessed_chs, selected_input_type, model, device)
+else:
+    dataset = create_raw_dataset(preprocessed_chs, device)

utils/pretraining.py

@@ -0,0 +1,150 @@
+#%% 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}")

utils/res1dcnn.py

@@ -0,0 +1,88 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.optim import Adam
+from torch.optim.lr_scheduler import MultiStepLR
+
+class ResidualBlock(nn.Module):
+
+    def __init__(self, in_channels, out_channels):
+        super(ResidualBlock, self).__init__()
+        self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size=3, padding=1)
+        self.bn1 = nn.BatchNorm1d(out_channels)
+        self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size=3, padding=1)
+        self.bn2 = nn.BatchNorm1d(out_channels)
+
+        # Shortcut connection to match dimensions when needed
+        self.shortcut = nn.Sequential()
+        if in_channels != out_channels:
+            self.shortcut = nn.Sequential(
+                nn.Conv1d(in_channels, out_channels, kernel_size=1),
+                nn.BatchNorm1d(out_channels)
+            )
+    
+    def forward(self, x):
+        residual = x
+        x = F.relu(self.bn1(self.conv1(x)))
+        x = self.bn2(self.conv2(x))
+        x += self.shortcut(residual)
+        x = F.relu(x)
+        return x
+
+
+class ResNet1DCNN(nn.Module):
+
+    def __init__(self, input_channels, sequence_length, num_classes):
+        super(ResNet1DCNN, self).__init__()
+        
+        # Initial convolution layer
+        self.conv1 = nn.Conv1d(input_channels, 32, kernel_size=7, stride=2, padding=3)
+        self.bn1 = nn.BatchNorm1d(32)
+        self.maxpool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)
+        
+        # Residual layers
+        self.layer1 = self._make_layer(32, 32, 2)
+        self.layer2 = self._make_layer(32, 64, 3)
+        self.layer3 = self._make_layer(64, 128, 4)
+        
+        # Calculate the size of the flattened features
+        with torch.no_grad():
+            dummy_input = torch.zeros(1, input_channels, sequence_length)
+            dummy_output = self.compute_conv_output(dummy_input)
+            self.flatten_size = dummy_output.numel()
+        
+        # Fully connected layers
+        self.fc1 = nn.Linear(self.flatten_size, 128)
+        self.bn_fc1 = nn.BatchNorm1d(128)
+        self.fc2 = nn.Linear(128, num_classes)
+        
+        self.dropout = nn.Dropout(0.5)
+        
+    def _make_layer(self, in_channels, out_channels, num_blocks):
+
+        layers = [ResidualBlock(in_channels, out_channels)]
+        for _ in range(1, num_blocks):
+            layers.append(ResidualBlock(out_channels, out_channels))
+        return nn.Sequential(*layers)
+    
+    def compute_conv_output(self, x):
+
+        x = self.maxpool(F.relu(self.bn1(self.conv1(x))))
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = F.adaptive_avg_pool1d(x, 1)
+        return x
+    
+    def forward(self, x):
+        
+        x = x.transpose(1, 2)
+        
+        x = self.compute_conv_output(x)
+        
+        x = x.view(x.size(0), -1)
+        x = F.relu(self.bn_fc1(self.fc1(x)))
+        x = self.dropout(x)
+        x = self.fc2(x)
+        
+        return x