Fix style errors in transformer code

felker · felker · commit 019c1bdf3ae8 · 2020-02-25T13:19:23.000-05:00
diff --git a/.travis.yml b/.travis.yml
@@ -2,6 +2,7 @@ language: python
 branches:
   only:
     - master
+    - transformer
 os:
   - linux
 
diff --git a/examples/transformer_learn.py b/examples/transformer_learn.py
@@ -75,13 +75,12 @@
     print('Training on {} shots, testing on {} shots'.format(
         len(shot_list_train), len(shot_list_test)))
 
-
     #####################################################
     #                    TRAINING                       #
     #####################################################
     train(conf, shot_list_train.random_sublist(512),
           shot_list_validate.random_sublist(256), loader)
-    #if not only_predict:
+    # if not only_predict:
     #    p = old_mp.Process(target=train,
     #                       args=(conf, shot_list_train,
     #                             shot_list_validate, loader)
@@ -115,18 +114,17 @@
 
     # TODO(KGF): check tuple unpack
     (y_prime_train, y_gold_train, disruptive_train, roc_train,
-    loss_train) = make_predictions_and_evaluate_gpu(
-        conf, shot_list_train, loader, custom_path)
+     loss_train) = make_predictions_and_evaluate_gpu(
+         conf, shot_list_train, loader, custom_path)
     (y_prime_test, y_gold_test, disruptive_test, roc_test,
-    loss_test) = make_predictions_and_evaluate_gpu(
-        conf, shot_list_test, loader, custom_path)
+     loss_test) = make_predictions_and_evaluate_gpu(
+         conf, shot_list_test, loader, custom_path)
     print('=========Summary========')
     print('Train Loss: {:.3e}'.format(loss_train))
     print('Train ROC: {:.4f}'.format(roc_train))
     print('Test Loss: {:.3e}'.format(loss_test))
     print('Test ROC: {:.4f}'.format(roc_test))
 
-
     disruptive_train = np.array(disruptive_train)
     disruptive_test = np.array(disruptive_test)
 
@@ -138,16 +136,20 @@
     shot_list_test.make_light()
     shot_list_train.make_light()
 
-    save_str = 'results_' + datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
+    save_str = 'results_' + datetime.datetime.now().strftime(
+        "%Y-%m-%d-%H-%M-%S")
     result_base_path = conf['paths']['results_prepath']
     if not os.path.exists(result_base_path):
         os.makedirs(result_base_path)
-    np.savez(result_base_path+save_str, y_gold=y_gold, y_gold_train=y_gold_train,
-            y_gold_test=y_gold_test, y_prime=y_prime, y_prime_train=y_prime_train,
-            y_prime_test=y_prime_test, disruptive=disruptive,
-            disruptive_train=disruptive_train, disruptive_test=disruptive_test,
-            shot_list_validate=shot_list_validate,
-            shot_list_train=shot_list_train, shot_list_test=shot_list_test,
-            conf=conf)
+    np.savez(result_base_path+save_str, y_gold=y_gold,
+             y_gold_train=y_gold_train,
+             y_gold_test=y_gold_test,
+             y_prime=y_prime, y_prime_train=y_prime_train,
+             y_prime_test=y_prime_test, disruptive=disruptive,
+             disruptive_train=disruptive_train,
+             disruptive_test=disruptive_test,
+             shot_list_validate=shot_list_validate,
+             shot_list_train=shot_list_train, shot_list_test=shot_list_test,
+             conf=conf)
 
     print('finished.')
diff --git a/plasma/models/distributed_torch_runner.py b/plasma/models/distributed_torch_runner.py
@@ -418,7 +418,8 @@ def train(conf, shot_list_train, shot_list_validate, loader):
     hvd.broadcast_parameters(train_model.state_dict(), root_rank=0)
     hvd.broadcast_optimizer_state(optim, root_rank=0)
 
-    optimizer_args = {'op': hvd.Average, 'compression': hvd.Compression.fp16, 'named_parameters': train_model.named_parameters()}
+    optimizer_args = {'op': hvd.Average, 'compression': hvd.Compression.fp16,
+                      'named_parameters': train_model.named_parameters()}
     optimizer = hvd.DistributedOptimizer(optim, **optimizer_args)
 
     train_model.train()
@@ -431,8 +432,8 @@ def train(conf, shot_list_train, shot_list_validate, loader):
     while e < num_epochs - 1:
         print('\nEpoch {}/{}'.format(e, num_epochs))
         (step, ave_loss, curr_loss, num_so_far,
-         effective_epochs) = train_epoch(train_model, data_gen, optimizer, scheduler,
-                                         loss_fn)
+         effective_epochs) = train_epoch(train_model, data_gen, optimizer,
+                                         scheduler, loss_fn)
         e = effective_epochs
         loader.verbose = False  # True during the first iteration
         # if task_index == 0:
diff --git a/plasma/models/loader.py b/plasma/models/loader.py
@@ -825,16 +825,15 @@ def get_batch_size(batch_size, prediction_mode):
     def get_num_skips(length, skip):
         return 1 + (length-1)//skip
 
-    #FIXME Alexeys
+    # FIXME Alexeys
     def simple_batch_generator(self, shot_list, max_len=2048, inference=False):
-
         batch_size = self.conf['training']['batch_size']
         sig, res = self.get_signal_result_from_shot(shot_list.shots[0])
         Xbuff = np.zeros((batch_size, max_len, sig.shape[1]))
         Ybuff = np.zeros((batch_size, max_len, res.shape[1]))
 
         num_total = len(shot_list)
-        #num_batches = num_total//batch_size
+        # num_batches = num_total//batch_size
         disr = np.zeros(batch_size, dtype=bool)
 
         while True:
@@ -846,18 +845,18 @@ def simple_batch_generator(self, shot_list, max_len=2048, inference=False):
             for i in range(num_total):
                 shot = self.sample_shot_from_list_given_index(shot_list, i)
                 sig, res = self.get_signal_result_from_shot(shot)
-                sig = sig[-max_len:,:]
-                res = res[-max_len:,:]
-                Xbuff[i%batch_size, -sig.shape[0]:, :] = sig
-                Ybuff[i%batch_size, -res.shape[0]:, :] = res
-                disr[i%batch_size] = shot.is_disruptive_shot()
+                sig = sig[-max_len:, :]
+                res = res[-max_len:, :]
+                Xbuff[i % batch_size, -sig.shape[0]:, :] = sig
+                Ybuff[i % batch_size, -res.shape[0]:, :] = res
+                disr[i % batch_size] = shot.is_disruptive_shot()
 
                 if i % batch_size == 0:
                     num_so_far += batch_size
-
                     yield Xbuff, Ybuff, num_so_far, num_total, disr
-                    #Xbuff = np.zeros((batch_size, max_len, sig.shape[1]))
-                    #Ybuff = np.zeros((batch_size, max_len, res.shape[1]))
+                    # Xbuff = np.zeros((batch_size, max_len, sig.shape[1]))
+                    # Ybuff = np.zeros((batch_size, max_len, res.shape[1]))
+
 
 class ProcessGenerator(object):
     def __init__(self, generator):
diff --git a/plasma/transformer/runner.py b/plasma/transformer/runner.py
@@ -8,17 +8,18 @@
 from plasma.utils.performance import PerformanceAnalyzer
 from plasma.utils.evaluation import get_loss_from_list
 from plasma.models.torch_runner import (
-    #make_predictions_and_evaluate_gpu,
-    #make_predictions,
+    # make_predictions_and_evaluate_gpu,
+    # make_predictions,
     get_signal_dimensions,
+    calculate_conv_output_size,
 )
 
 from functools import partial
 import os
 import numpy as np
 import logging
 import random
-import tqdm
+# import tqdm
 
 model_filename = "torch_model.pt"
 LOGGER = logging.getLogger("plasma.transformer.runner")
@@ -32,12 +33,14 @@
 # else:
 #    device = torch.device("cpu")
 
+
 def set_seed(seed):
     random.seed(seed)
     np.random.seed(seed)
     torch.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)
-    os.environ["PYTHONHASHSEED"]="0"
+    os.environ["PYTHONHASHSEED"] = "0"
+
 
 class TransformerNet(nn.Module):
     def __init__(
@@ -103,14 +106,16 @@ def __init__(
                 )
                 self.layers.append(nn.MaxPool1d(kernel_size=self.pooling_size))
                 self.conv_output_size = calculate_conv_output_size(
-                    self.conv_output_size, 0, 1, self.pooling_size, self.pooling_size
+                    self.conv_output_size, 0, 1, self.pooling_size,
+                    self.pooling_size
                 )
                 self.layers.append(nn.Dropout2d(dropout))
             self.net = nn.Sequential(*self.layers)
             self.conv_output_size = self.conv_output_size * layer_sizes[-1]
         self.linear_layers = []
 
-        print("Final feature size = {}".format(self.n_scalars + self.conv_output_size))
+        print("Final feature size = {}".format(self.n_scalars
+                                               + self.conv_output_size))
         self.linear_layers.append(
             nn.Linear(self.conv_output_size + self.n_scalars, linear_size)
         )
@@ -128,7 +133,7 @@ def forward(self, x):
                 x_profiles = x
             else:
                 x_scalars = x[:, : self.n_scalars]
-                x_profiles = x[:, self.n_scalars :]
+                x_profiles = x[:, self.n_scalars:]
             x_profiles = x_profiles.contiguous().view(
                 x.size(0), self.n_profiles, self.profile_size
             )
@@ -170,7 +175,8 @@ def __init__(
         self.__max_seq_length = max_seq_length
         self.__d_model = d_model
         # FIXME
-        self.__positional_encodings = nn.Embedding(max_seq_length, d_model).float()
+        self.__positional_encodings = nn.Embedding(
+            max_seq_length, d_model).float()
 
     def forward(self, x):
         """
@@ -180,17 +186,19 @@ def forward(self, x):
         mask = (
             torch.arange(x.shape[1], device=device)
             .unsqueeze(0)
-            .lt(torch.tensor([self.__max_seq_length], device=device).unsqueeze(-1))
+            .lt(torch.tensor([self.__max_seq_length],
+                             device=device).unsqueeze(-1))
         )
         transformer_input = x * mask.unsqueeze(-1).float()  # B x max_len x D
 
         positional_encodings = self.__positional_encodings(
             torch.arange(x.shape[1], dtype=torch.int64, device=device)
         ).unsqueeze(0)
-        transformer_input = transformer_input + positional_encodings  # B x max_len x D
+        transformer_input = (transformer_input
+                             + positional_encodings)  # B x max_len x D
 
         out = self.__transformer_encoder(
-            transformer_input #.transpose(0, 1), src_key_padding_mask=~mask
+            transformer_input  # .transpose(0, 1), src_key_padding_mask=~mask
         )
         return out
 
@@ -199,11 +207,10 @@ def build_torch_model(conf):
 
     dropout = conf["model"]["dropout_prob"]
     n_scalars, n_profiles, profile_size = get_signal_dimensions(conf)
-
-    output_size = 1
-    layer_sizes_spatial = [6, 3, 3] 
+    # output_size = 1
+    layer_sizes_spatial = [6, 3, 3]
     kernel_size_spatial = 3
-    linear_size = 5 #FIXME Alexeys there will be no linear layers
+    linear_size = 5  # FIXME Alexeys there will be no linear layers
 
     model = TransformerNet(
         n_scalars,
@@ -233,7 +240,7 @@ def train_epoch(model, data_gen, optimizer, scheduler, loss_fn):
     step = 0
     while True:
         x_, y_, num_so_far, num_total, _ = next(data_gen)
-      
+
         x = torch.from_numpy(x_).float().to(device)
         y = torch.from_numpy(y_).float().to(device)
 
@@ -247,75 +254,77 @@ def train_epoch(model, data_gen, optimizer, scheduler, loss_fn):
         scheduler.step()
         step += 1
 
-        LOGGER.info(
-            f"[{step}]  [{num_so_far}/{num_total}] loss: {loss.item()}, ave_loss: {total_loss / step}"
-            )
+        LOGGER.info(f"[{step}]  [{num_so_far}/{num_total}] loss: {loss.item()}, ave_loss: {total_loss / step}")  # noqa
         if num_so_far >= num_total:
             break
 
-    return step, loss.item(), total_loss, num_so_far, 1.0 * num_so_far / num_total
+    return (step, loss.item(), total_loss, num_so_far,
+            1.0 * num_so_far / num_total)
 
 
 def train(conf, shot_list_train, shot_list_validate, loader):
-    #set random seed
+    # set random seed
     set_seed(0)
     num_epochs = conf["training"]["num_epochs"]
-    patience = conf["callbacks"]["patience"]
+    # patience = conf["callbacks"]["patience"]
     lr_decay = conf["model"]["lr_decay"]
-    batch_size = conf['training']['batch_size']
+    # batch_size = conf['training']['batch_size']
     lr = conf["model"]["lr"]
-    clipnorm = conf['model']['clipnorm']
+    # clipnorm = conf['model']['clipnorm']
     e = 0
 
     loader.set_inference_mode(False)
     train_data_gen = partial(
         loader.simple_batch_generator,
         shot_list=shot_list_train,
     )()
-    valid_data_generator = partial(
+    valid_data_generator = partial(  # noqa
         loader.simple_batch_generator,
         shot_list=shot_list_validate,
         inference=True
     )()
-    LOGGER.info(f"validate: {len(shot_list_validate)} shots, {shot_list_validate.num_disruptive()} disruptive")
-    LOGGER.info(f"training: {len(shot_list_train)} shots, {shot_list_train.num_disruptive()} disruptive")
+    LOGGER.info(f"validate: {len(shot_list_validate)} shots, {shot_list_validate.num_disruptive()} disruptive")  # noqa
+    LOGGER.info(f"training: {len(shot_list_train)} shots, {shot_list_train.num_disruptive()} disruptive")  # noqa
 
     loss_fn = nn.MSELoss(size_average=True)
     train_model = build_torch_model(conf)
 
     optimizer = opt.Adam(train_model.parameters(), lr=lr)
     scheduler = opt.lr_scheduler.ExponentialLR(optimizer, lr_decay)
- 
+
     model_path = get_model_path(conf)
     makedirs_process_safe(os.path.dirname(model_path))
 
     train_model.train()
     LOGGER.info(f"{num_epochs - 1 - e} epochs left to go")
     while e < num_epochs - 1:
         LOGGER.info(f"Epoch {e}/{num_epochs}")
-        (step, ave_loss, curr_loss, num_so_far, effective_epochs) = train_epoch(
+        (step, ave_loss, curr_loss, num_so_far,
+         effective_epochs) = train_epoch(
             train_model, train_data_gen, optimizer, scheduler, loss_fn
         )
-        
+
         e = effective_epochs
         torch.save(train_model.state_dict(), model_path)
-        #FIXME no validation for now as OOM
-        #_, _, _, roc_area, loss = make_predictions_and_evaluate_gpu(
+        # FIXME no validation for now as OOM
+        # _, _, _, roc_area, loss = make_predictions_and_evaluate_gpu(
         #    conf, shot_list_validate, valid_data_generator
-        #)
+        # )
+
+        # # stop_training = False
+        # print("=========Summary======== for epoch{}".format(step))
+        # print("Training Loss numpy: {:.3e}".format(ave_loss))
+        # print("Validation Loss: {:.3e}".format(loss))
+        # print("Validation ROC: {:.4f}".format(roc_area))
 
-        ## stop_training = False
-        #print("=========Summary======== for epoch{}".format(step))
-        #print("Training Loss numpy: {:.3e}".format(ave_loss))
-        #print("Validation Loss: {:.3e}".format(loss))
-        #print("Validation ROC: {:.4f}".format(roc_area))
 
 def apply_model_to_np(model, x):
     return model(torch.from_numpy(x).float()).data.numpy()
 
-#FIXME Alexeys change
+
+# FIXME Alexeys change
 def make_predictions(conf, shot_list, generator, custom_path=None):
-    #generator = loader.inference_batch_generator_full_shot(shot_list)
+    # generator = loader.inference_batch_generator_full_shot(shot_list)
     inference_model = build_torch_model(conf)
 
     if custom_path is None:
@@ -336,11 +345,11 @@ def make_predictions(conf, shot_list, generator, custom_path=None):
 
         x = torch.from_numpy(x_).float().to(device)
         y = torch.from_numpy(y_).float().to(device)
-        #output = apply_model_to_np(inference_model, x)
+        # output = apply_model_to_np(inference_model, x)
         output = inference_model(x)
 
         for batch_idx in range(x.shape[0]):
-            #curr_length = lengths[batch_idx]
+            # curr_length = lengths[batch_idx]
             y_prime += [output[batch_idx, :, 0]]
             y_gold += [y[batch_idx, :, 0]]
             disruptive += [disr[batch_idx]]
@@ -352,11 +361,13 @@ def make_predictions(conf, shot_list, generator, custom_path=None):
             break
     return y_prime, y_gold, disruptive
 
-#FIXME ALexeys change loader --> generator
-def make_predictions_and_evaluate_gpu(conf, shot_list, generator, custom_path=None):
+
+# FIXME ALexeys change loader --> generator
+def make_predictions_and_evaluate_gpu(conf, shot_list, generator,
+                                      custom_path=None):
     y_prime, y_gold, disruptive = make_predictions(
         conf, shot_list, generator, custom_path)
     analyzer = PerformanceAnalyzer(conf=conf)
     roc_area = analyzer.get_roc_area(y_prime, y_gold, disruptive)
     loss = get_loss_from_list(y_prime, y_gold, conf['data']['target'])
-    return y_prime, y_gold, disruptive, roc_area, loss
+    return y_prime, y_gold, disruptive, roc_area, loss
