Upload tutorial.py

tutorial.py

@@ -74,11 +74,13 @@ scenario_names = np.array([
 scenario_idxs = np.array([0, 1, 2, 3, 4, 5])[3] 
 selected_scenario_names = scenario_names[scenario_idxs]
 
+snr_db = None
+
 preprocessed_chs = tokenizer(
     selected_scenario_names=selected_scenario_names,
     manual_data=None,
     gen_raw=True,
-    snr_db=None
+    snr_db=snr_db
 )
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -87,7 +89,7 @@ model = lwm.from_pretrained(device=device)
 #%%
 from inference import lwm_inference, create_raw_dataset
 input_types = ['cls_emb', 'channel_emb', 'raw']
-selected_input_type = input_types[2] 
+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)
@@ -139,7 +141,7 @@ for task in tasks:
 #%% TRAINING
 #%% TRAINING PARAMETERS
 task = ['LoS/NLoS Classification', 'Beam Prediction'][0] # Select the task
-n_trials = 10  # Number of trials for each configuration
+n_trials = 1  # Number of trials for each configuration
 num_classes = 2 if task == 'LoS/NLoS Classification' else n_beams  # Set number of classes based on the task
 input_types = ['raw', 'cls_emb']  # Types of input data
 split_ratios = np.array([.005, .0075, .01, .015, .02, .03, 
@@ -174,6 +176,20 @@ for input_type_idx, input_type in enumerate(input_types):
             print(f"\ninput type: {input_type}, \nnumber of training samples: {int(split_ratio*len(dataset))}, \ntrial: {trial}\n")
             
             torch.manual_seed(trial)  # Set seed for reproducibility
+            
+            if snr_db is not None:
+                preprocessed_chs = tokenizer(
+                    selected_scenario_names=selected_scenario_names,
+                    manual_data=None,
+                    gen_raw=True,
+                    snr_db=snr_db
+                )
+                if input_type in ['cls_emb', 'channel_emb']:
+                    dataset = lwm_inference(preprocessed_chs, input_type, model, device)
+                else:
+                    dataset = create_raw_dataset(preprocessed_chs, device)
+                dataset = dataset.view(dataset.size(0), -1)
+                
             train_loader, test_loader = get_data_loaders(
                 dataset, 
                 labels, 
@@ -240,6 +256,20 @@ for input_type_idx, input_type in enumerate(input_types):
         for trial in range(n_trials):
             
             torch.manual_seed(trial)  # Set seed for reproducibility
+            
+            if snr_db is not None:
+                preprocessed_chs = tokenizer(
+                    selected_scenario_names=selected_scenario_names,
+                    manual_data=None,
+                    gen_raw=True,
+                    snr_db=snr_db
+                )
+                if input_type in ['cls_emb', 'channel_emb']:
+                    dataset = lwm_inference(preprocessed_chs, input_type, model, device)
+                else:
+                    dataset = create_raw_dataset(preprocessed_chs, device)
+                dataset = dataset.view(dataset.size(0), -1)
+            
             train_loader, test_loader = get_data_loaders(
                 dataset, 
                 labels,