Merge pull request #37 from clessig/iluise/head

move to Dask array + first working version of the multiformer

Author: clessig

atmorep/core/atmorep_model.py

@@ -17,6 +17,7 @@
 import torch
 import numpy as np
 import code
+import os
 # code.interact(local=locals())
 
 # import horovod.torch as hvd
@@ -238,8 +239,18 @@ def create( self, devices, load_pretrained=True) :
       if len(field_info[1]) > 4 and load_pretrained :
         # TODO: inconsistent with embeds_token_info -> version that can handle both
         #       we could imply use the file name: embed_token_info vs embeds_token_info
-        name = 'AtmoRep' + '_embed_token_info'
+        name = 'AtmoRep' + '_embeds_token_info'
+        if not os.path.exists(get_model_filename( name, field_info[1][4][0], field_info[1][4][1])):
+          name = 'AtmoRep' + '_embed_token_info'
+
         mloaded = torch.load( get_model_filename( name, field_info[1][4][0], field_info[1][4][1]))
+      
+        if "weight" not in mloaded.keys(): #TODO: get rid of this
+          mloaded["weight"] = mloaded["0.weight"]
+          mloaded["bias"] = mloaded["0.bias"]
+          del mloaded["0.weight"]
+          del mloaded["0.bias"]
+
         self.embeds_token_info[-1].load_state_dict( mloaded)
         print( 'Loaded embed_token_info from id = {}.'.format( field_info[1][4][0] ) )
       else :
@@ -260,7 +271,7 @@ def create( self, devices, load_pretrained=True) :
       if len(field_info[1]) > 4 and load_pretrained :
         self.load_block( field_info, 'encoder', self.encoders[-1])
       self.embeds.append( self.encoders[-1].embed)
-
+     
       # indices of coupled fields for efficient access in forward
       self.fields_coupling_idx.append( [field_idx])
       for field_coupled in field_info[1][2] : 
@@ -353,9 +364,17 @@ def load_block( self, field_info, block_name, block ) :
                 param[ : , b_loaded[k].shape[1] : ] = 0.01 * torch.rand( param.shape[0],
                                                                          param.shape[1] - b_loaded[k].shape[1])
                 keys_del += [ k ]
+
+      #for backward compatibility. solved in new runs
+      if 'proj_heads_other' in name: 
+        for k in b_loaded.keys() :
+          if name == k :
+            if b_loaded[name].shape[0] == 0:
+              keys_del += [ name ]
+      
     for k in keys_del :
       del b_loaded[k]
-
+    
     # use strict=False so that differing blocks, e.g. through coupling, are ignored
     mkeys, _ = block.load_state_dict( b_loaded, False)
 
@@ -394,11 +413,7 @@ def translate_weights(self, mloaded, mkeys, ukeys):
           del mloaded[f'encoders.0.heads.{layer}.heads_other.{head}.proj_ks.weight']
           del mloaded[f'encoders.0.heads.{layer}.heads_other.{head}.proj_vs.weight']
 
-      else:
-        dim_mw = self.encoders[0].heads[0].proj_heads_other[0].weight.shape
-        mw = torch.tensor(np.zeros(dim_mw))
-      
-      mloaded[f'encoders.0.heads.{layer}.proj_heads_other.0.weight'] = mw  
+        mloaded[f'encoders.0.heads.{layer}.proj_heads_other.0.weight'] = mw  
 
     #decoder
     for iblock in range(0, 19, 2) : 

atmorep/core/evaluate.py

@@ -19,24 +19,24 @@
 
 if __name__ == '__main__':
 
-  # models for individual fields
+  # arXiv 2023: models for individual fields
   #model_id = '4nvwbetz'     # vorticity
   #model_id = 'oxpycr7w'     # divergence
   #model_id = '1565pb1f'     # specific_humidity
   #model_id = '3kdutwqb'     # total precip
-  model_id = 'dys79lgw'     # velocity_u
+  #model_id = 'dys79lgw'     # velocity_u
   #model_id = '22j6gysw'     # velocity_v
-  # model_id = '15oisw8d'     # velocity_z
-  #model_id = '3qou60es'     # temperature (also 2147fkco)
+  #model_id = '15oisw8d'     # velocity_z
+  #model_id = '3qou60es'     # temperature 
   #model_id = '2147fkco'     # temperature (also 2147fkco)
- 
-  # multi-field configurations with either velocity or voritcity+divergence
-  #model_id = '1jh2qvrx'     # multiformer, velocity
-  # model_id = 'wqqy94oa'     # multiformer, vorticity
-  #model_id = '3cizyl1q'     # 3 field config: u,v,T
-  # model_id = '1v4qk0qx'     # pre-trained, 3h forecasting
-  # model_id = '1m79039j'     # pre-trained, 6h forecasting
-  #model_id='34niv2nu'
+  
+  # new runs 2024
+  #model_id='j8dwr5qj' #velocity_u
+  #model_id='0tlnm5up' #velocity_v
+  #model_id='v63l01zu' #specific humidity 
+  #model_id='9l1errbo' #velocity_z
+  model_id='7ojls62c' #temperature 1024 
+  
   # supported modes: test, forecast, fixed_location, temporal_interpolation, global_forecast,
   #                  global_forecast_range
   # options can be used to over-write parameters in config; some modes also have specific options, 
@@ -45,7 +45,7 @@
   #Add 'attention' : True to options to store the attention maps. NB. supported only for single field runs. 
 
   # BERT masked token model
-  mode, options = 'BERT', {'years_test' : [2021], 'num_samples_validate' : 128, 'with_pytest' : True }
+  #mode, options = 'BERT', {'years_val' : [2021], 'num_samples_validate' : 128, 'with_pytest' : True}
 
   # BERT forecast mode
   #mode, options = 'forecast', {'forecast_num_tokens' : 2, 'num_samples_validate' : 128, 'with_pytest' : True }
@@ -55,19 +55,19 @@
   #mode, options = 'temporal_interpolation', {'idx_time_mask': [5,6,7], 'num_samples_validate' : 128, 'with_pytest' : True}
 
   # BERT forecast with patching to obtain global forecast
-#   mode, options = 'global_forecast', { 
-#                                       'dates' : [[2021, 1, 10, 18]],
-# #                                     #   'dates' : [ #[2021, 1, 10, 18]
-# #                                     #      [2021, 1, 10, 12] , [2021, 1, 11, 0], [2021, 1, 11, 12], [2021, 1, 12, 0], [2021, 1, 12, 12], [2021, 1, 13, 0], 
-# #                                     #      [2021, 4, 10, 12], [2021, 4, 11, 0], [2021, 4, 11, 12], [2021, 4, 12, 0], [2021, 4, 12, 12], [2021, 4, 13, 0], 
-# #                                     #      [2021, 7, 10, 12], [2021, 7, 11, 0], [2021, 7, 11, 12], [2021, 7, 12, 0], [2021, 7, 12, 12], [2021, 7, 13, 0], 
-# #                                     #      [2021, 10, 10, 12], [2021, 10, 11, 0], [2021, 10, 11, 12], [2021, 10, 12, 0], [2021, 10, 12, 12], [2021, 10, 13, 0]
-# #                                     #  ], 
-#                                       'token_overlap' : [0, 0],
-#                                       'forecast_num_tokens' : 2,
-#                                       'with_pytest' : True }
+  mode, options = 'global_forecast', { 
+                                      #'dates' : [[2021, 2, 10, 12]]
+                                      'dates' : [
+                                         [2021, 1, 10, 12] , [2021, 1, 11, 0], [2021, 1, 11, 12], [2021, 1, 12, 0], #[2021, 1, 12, 12], [2021, 1, 13, 0], 
+                                         [2021, 4, 10, 12], [2021, 4, 11, 0], [2021, 4, 11, 12], [2021, 4, 12, 0], #[2021, 4, 12, 12], [2021, 4, 13, 0], 
+                                         [2021, 7, 10, 12], [2021, 7, 11, 0], [2021, 7, 11, 12], [2021, 7, 12, 0], #[2021, 7, 12, 12], [2021, 7, 13, 0], 
+                                         [2021, 10, 10, 12], [2021, 10, 11, 0], [2021, 10, 11, 12], #[2021, 10, 12, 0], [2021, 10, 12, 12], [2021, 10, 13, 0]
+                                       ], 
+                                      'token_overlap' : [0, 0],
+                                      'forecast_num_tokens' : 2,
+                                      'with_pytest' : True }
 
-  file_path = '/gpfs/scratch/ehpc03/era5_y2010_2021_res025_chunk8.zarr'
+  file_path = '/gpfs/scratch/ehpc03/era5_y1979_2021_res025_chunk8.zarr'
 
   now = time.time()
   Evaluator.evaluate( mode, model_id, file_path, options)

atmorep/core/evaluator.py

@@ -69,18 +69,13 @@ def run( cf, model_id, model_epoch, devices) :
   @staticmethod
   def evaluate( mode, model_id, file_path, args = {}, model_epoch=-2) :
 
-    # SLURM_TASKS_PER_NODE is controlled by #SBATCH --ntasks-per-node=1; should be 1 for multiformer
-    with_ddp = True
-    if '-1' == os.environ.get('MASTER_ADDR', '-1') :
-      with_ddp = False
-      num_accs_per_task = 1 
-    else :
-      num_accs_per_task = int( 4 / int( os.environ.get('SLURM_TASKS_PER_NODE', '1')[0] ))
-    devices = init_torch( num_accs_per_task)
-    #devices = ['cuda:1']
-
+    devices = init_torch()
+    with_ddp = True 
     par_rank, par_size = setup_ddp( with_ddp)
+    
     cf = Config().load_json( model_id)
+  
+    cf.num_accs_per_task = len(devices)
     cf.file_path = file_path
     cf.with_wandb = True
     cf.with_ddp = with_ddp
@@ -120,7 +115,7 @@ def evaluate( mode, model_id, file_path, args = {}, model_epoch=-2) :
     if cf.with_pytest:
       fields = [field[0] for field in cf.fields_prediction]
       for field in fields:
-        pytest.main(["-x", "./atmorep/tests/validation_test.py", "--field", field, "--model_id", cf.wandb_id, "--strategy", cf.BERT_strategy])
+        pytest.main(["-x", "-s", "./atmorep/tests/validation_test.py", "--field", field, "--model_id", cf.wandb_id, "--strategy", cf.BERT_strategy])
 
   ##############################################
   @staticmethod
@@ -155,11 +150,11 @@ def global_forecast( cf, model_id, model_epoch, devices, args = {}) :
     cf.batch_size_test = 24
     cf.num_loader_workers = 12 #1
     cf.log_test_num_ranks = 1
+
+    #TODO: temporary solution. Add support for batch_size > 1 
+    cf.batch_size_validation = 1 #64
+    cf.batch_size = 1
 
-    if not hasattr(cf, 'batch_size'):
-      cf.batch_size = 196 #14
-    if not hasattr(cf, 'batch_size_validation'):
-      cf.batch_size_validation = 1 #64
     if not hasattr(cf, 'num_samples_validate'):
       cf.num_samples_validate = 196 
     #if not hasattr(cf,'with_mixed_precision'):

atmorep/core/train.py

@@ -31,13 +31,13 @@
 ####################################################################################################
 def train_continue( wandb_id, epoch, Trainer, epoch_continue = -1) :
 
-  num_accs_per_task = int( 4 / int( os.environ.get('SLURM_TASKS_PER_NODE', '1')[0] ))
-  device = init_torch( num_accs_per_task)
+  devices = init_torch()
   with_ddp = True
   par_rank, par_size = setup_ddp( with_ddp)
 
   cf = Config().load_json( wandb_id)
 
+  cf.num_accs_per_task = len(devices)   # number of GPUs / accelerators per task
   cf.with_ddp = with_ddp
   cf.par_rank = par_rank
   cf.par_size = par_size
@@ -56,13 +56,13 @@ def train_continue( wandb_id, epoch, Trainer, epoch_continue = -1) :
     cf.with_mixed_precision = True
   if not hasattr(cf, 'years_val'):
     cf.years_val = cf.years_test
-  
+ 
   # any parameter in cf can be overwritten when training is continued, e.g. we can increase the 
   # masking rate 
   # cf.fields = [ [ 'specific_humidity', [ 1, 2048, [ ], 0 ], 
   #                               [ 96, 105, 114, 123, 137 ], 
   #                               [12, 6, 12], [3, 9, 9], [0.5, 0.9, 0.1, 0.05] ] ]
-
+  
   setup_wandb( cf.with_wandb, cf, par_rank, project_name='train', mode='offline')  
   # resuming a run requires online mode, which is not available everywhere
   #setup_wandb( cf.with_wandb, cf, par_rank, wandb_id = wandb_id)  
@@ -75,15 +75,14 @@ def train_continue( wandb_id, epoch, Trainer, epoch_continue = -1) :
     epoch_continue = epoch
 
   # run
-  trainer = Trainer.load( cf, wandb_id, epoch, device)
+  trainer = Trainer.load( cf, wandb_id, epoch, devices)
   print( 'Loaded run \'{}\' at epoch {}.'.format( wandb_id, epoch))
   trainer.run( epoch_continue)
 
 ####################################################################################################
 def train() :
 
-  num_accs_per_task = int( 4 / int( os.environ.get('SLURM_TASKS_PER_NODE', '1')[0] ))
-  device = init_torch( num_accs_per_task)
+  devices = init_torch()
   with_ddp = True
   par_rank, par_size = setup_ddp( with_ddp)
 
@@ -93,7 +92,7 @@ def train() :
   cf = Config()
   # parallelization
   cf.with_ddp = with_ddp
-  cf.num_accs_per_task = num_accs_per_task   # number of GPUs / accelerators per task
+  cf.num_accs_per_task = len(devices)   # number of GPUs / accelerators per task
   cf.par_rank = par_rank
   cf.par_size = par_size
 
@@ -108,32 +107,31 @@ def train() :
 
   # cf.fields = [ [ 'temperature', [ 1, 1024, [ ], 0 ], 
   #                              [ 96, 105, 114, 123, 137 ], 
-  #                              [12, 6, 12], [3, 9, 9], [0.25, 0.9, 0.1, 0.05], 'local' ] ]
+  #                              [12, 2, 4], [3, 27, 27], [0.5, 0.9, 0.2, 0.05], 'local' ] ]
   # cf.fields_prediction = [ [cf.fields[0][0], 1.] ]
-
+ 
   cf.fields = [ [ 'velocity_u', [ 1, 1024, [ ], 0 ], 
                                 [ 96, 105, 114, 123, 137 ], 
                                  [12, 3, 6], [3, 18, 18], [0.5, 0.9, 0.2, 0.05] ] ]
 
   cf.fields_prediction = [ [cf.fields[0][0], 1.] ]
 
-
+  
   # cf.fields = [ [ 'velocity_v', [ 1, 1024, [ ], 0 ], 
   #                               [ 96, 105, 114, 123, 137 ], 
   #                               [12, 3, 6], [3, 18, 18], [0.25, 0.9, 0.1, 0.05] ] ]
 
-  # cf.fields = [ [ 'velocity_z', [ 1, 512, [ ], 0 ], 
+  # cf.fields = [ [ 'velocity_z', [ 1, 1024, [ ], 0 ], 
   #                               [ 96, 105, 114, 123, 137 ], 
-  #                               [12, 6, 12], [3, 9, 9], [0.25, 0.9, 0.1, 0.05] ] ]
+  #                               [12, 3, 6], [3, 18, 18], [0.25, 0.9, 0.1, 0.05] ] ]
 
   # cf.fields = [ [ 'specific_humidity', [ 1, 1024, [ ], 0 ], 
   #                               [ 96, 105, 114, 123, 137 ], 
-  #                               [12, 6, 12], [3, 9, 9], [0.25, 0.9, 0.1, 0.05], 'local' ] ]
-  #                             [12, 2, 4], [3, 27, 27], [0.5, 0.9, 0.1, 0.05], 'local' ] ]
-
+  #                               [12, 3, 6], [3, 18, 18], [0.25, 0.9, 0.1, 0.05] ] ]
+  
   cf.fields_targets = []
 
-  cf.years_train = list( range( 2010, 2021))
+  cf.years_train = list( range( 1979, 2021))
   cf.years_val = [2021]  #[2018] 
   cf.month = None
   cf.geo_range_sampling = [[ -90., 90.], [ 0., 360.]]
@@ -143,7 +141,7 @@ def train() :
   # training params
   cf.batch_size_validation = 1 #64
   cf.batch_size = 96
-  cf.num_epochs = 128
+  cf.num_epochs = 400 #128
   cf.num_samples_per_epoch = 4096*12
   cf.num_samples_validate = 128*12
   cf.num_loader_workers = 8
@@ -176,7 +174,7 @@ def train() :
   cf.net_tail_num_nets = 16
   cf.net_tail_num_layers = 0
   # loss
-  cf.losses = ['mse_ensemble', 'stats']  # mse, mse_ensemble, stats, crps, weighted_mse
+  cf.losses = ['mse_ensemble', 'stats'] # mse, mse_ensemble, stats, crps, weighted_mse
   # training
   cf.optimizer_zero = False
   cf.lr_start = 5. * 10e-7
@@ -185,6 +183,7 @@ def train() :
   cf.weight_decay = 0.05 #0.1
   cf.lr_decay_rate = 1.025
   cf.lr_start_epochs = 3
+  cf.model_log_frequency = 256 #save checkpoint every X batches
   # BERT
   # strategies: 'BERT', 'forecast', 'temporal_interpolation'
   cf.BERT_strategy = 'BERT'
@@ -218,7 +217,7 @@ def train() :
   # # # cf.file_path = '/p/scratch/atmo-rep/data/era5_1deg/months/era5_y2021_res025_chunk8.zarr'
   # # cf.file_path = '/ec/res4/scratch/nacl/atmorep/era5_y2021_res025_chunk8_lat180_lon180.zarr'
   # # # cf.file_path = '/ec/res4/scratch/nacl/atmorep/era5_y2021_res025_chunk16.zarr'
-  cf.file_path = '/gpfs/scratch/ehpc03/era5_y2010_2021_res025_chunk8.zarr/'
+  cf.file_path = '/gpfs/scratch/ehpc03/era5_y1979_2021_res025_chunk8.zarr/'
   # # # in steps x lat_degrees x lon_degrees
   cf.n_size = [36, 0.25*9*6, 0.25*9*12]
 
@@ -230,7 +229,7 @@ def train() :
     cf.write_json( wandb)
     cf.print()
 
-  trainer = Trainer_BERT( cf, device).create()
+  trainer = Trainer_BERT( cf, devices).create()
   trainer.run()
 
 ####################################################################################################
@@ -239,8 +238,8 @@ def train() :
   try :
 
     train()
-
-    #  wandb_id, epoch, epoch_continue = '1jh2qvrx', 392, 392
+ 
+    #  wandb_id, epoch, epoch_continue = 'gxfywjzl', 127, 127
     #  Trainer = Trainer_BERT
     #  train_continue( wandb_id, epoch, Trainer, epoch_continue)