Lint algoperf/

Author: fsschneider

algoperf/checkpoint_utils.py

@@ -7,14 +7,14 @@
 import os
 from typing import Sequence, Tuple
 
+import jax
+import numpy as np
+import torch
 from absl import logging
 from flax import jax_utils
 from flax.training import checkpoints as flax_checkpoints
 from flax.training.checkpoints import latest_checkpoint
-import jax
-import numpy as np
 from tensorflow.io import gfile  # pytype: disable=import-error
-import torch
 
 from algoperf import spec
 from algoperf.pytorch_utils import pytorch_setup

algoperf/data_utils.py

@@ -7,9 +7,7 @@
 import torch
 import torch.distributed as dist
 import torch.nn.functional as F
-from torch.utils.data import DataLoader
-from torch.utils.data import DistributedSampler
-from torch.utils.data import Sampler
+from torch.utils.data import DataLoader, DistributedSampler, Sampler
 
 from algoperf import spec
 
@@ -266,7 +264,7 @@ def prefetched_loader(self) -> Iterable[Tuple[spec.Tensor, spec.Tensor]]:
         next_targets = next_targets.to(self.device, non_blocking=True)
 
       if not first:
-        yield inputs, targets
+        yield inputs, targets  # noqa: F821
       else:
         first = False
 

algoperf/logger_utils.py

@@ -11,12 +11,12 @@
 import sys
 from typing import Any, Dict, Optional
 
-from absl import flags
-from clu import metric_writers
 import GPUtil
 import pandas as pd
 import psutil
 import torch.distributed as dist
+from absl import flags
+from clu import metric_writers
 
 from algoperf import spec
 from algoperf.pytorch_utils import pytorch_setup
@@ -198,7 +198,7 @@ def _get_system_hardware_info() -> Dict:
   try:
     system_hardware_info['cpu_model_name'] = _get_cpu_model_name()
     system_hardware_info['cpu_count'] = psutil.cpu_count()
-  except:  # pylint: disable=bare-except
+  except:  # noqa: E722
     logging.info('Unable to record cpu information. Continuing without it.')
 
   gpus = GPUtil.getGPUs()
@@ -207,7 +207,7 @@ def _get_system_hardware_info() -> Dict:
       system_hardware_info['gpu_model_name'] = gpus[0].name
       system_hardware_info['gpu_count'] = len(gpus)
       system_hardware_info['gpu_driver'] = gpus[0].driver
-    except:  # pylint: disable=bare-except
+    except:  # noqa: E722
       logging.info('Unable to record gpu information. Continuing without it.')
 
   return system_hardware_info
@@ -232,7 +232,7 @@ def _get_system_software_info() -> Dict:
     system_software_info['git_commit_hash'] = _get_git_commit_hash()
     # Note: do not store git repo url as it may be sensitive or contain a
     # secret.
-  except:  # pylint: disable=bare-except
+  except:  # noqa: E722
     logging.info('Unable to record git information. Continuing without it.')
 
   return system_software_info
@@ -279,7 +279,7 @@ def _get_workload_properties(workload: spec.Workload) -> Dict:
   for key in keys:
     try:
       attr = getattr(workload, key)
-    except:  # pylint: disable=bare-except
+    except:  # noqa: E722
       logging.info(
         f'Unable to record workload.{key} information. Continuing without it.'
       )

algoperf/profiler.py

@@ -4,10 +4,10 @@
 https://github.com/Lightning-AI/lightning/tree/master/src/pytorch_lightning/profilers.
 """
 
-from collections import defaultdict
-from contextlib import contextmanager
 import os
 import time
+from collections import defaultdict
+from contextlib import contextmanager
 from typing import Dict, Generator, List, Optional, Tuple
 
 import numpy as np

algoperf/pytorch_utils.py

@@ -1,11 +1,11 @@
 import os
 from typing import Tuple
 
-from absl import logging
 import jax
 import tensorflow as tf
 import torch
 import torch.distributed as dist
+from absl import logging
 
 from algoperf import spec
 from algoperf.profiler import Profiler

algoperf/random_utils.py

@@ -2,9 +2,8 @@
 
 from typing import Any, List, Union
 
-from absl import flags
-from absl import logging
 import numpy as np
+from absl import flags, logging
 
 try:
   import jax.random as jax_rng

algoperf/spec.py

@@ -5,10 +5,10 @@
 import functools
 from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union
 
-from absl import logging
 import jax
-from torch import nn
 import torch.nn.functional as F
+from absl import logging
+from torch import nn
 
 
 class LossType(enum.Enum):

algoperf/workloads/cifar/cifar_jax/input_pipeline.py

@@ -8,10 +8,10 @@
 import functools
 from typing import Dict, Iterator, Tuple
 
-from flax import jax_utils
 import jax
 import tensorflow as tf
 import tensorflow_datasets as tfds
+from flax import jax_utils
 
 from algoperf import spec
 from algoperf.data_utils import shard_and_maybe_pad_np

algoperf/workloads/cifar/cifar_jax/models.py

@@ -7,8 +7,8 @@
 import functools
 from typing import Any, Callable, Tuple
 
-from flax import linen as nn
 import jax.numpy as jnp
+from flax import linen as nn
 
 from algoperf import spec
 from algoperf.workloads.imagenet_resnet.imagenet_jax.models import ResNetBlock

algoperf/workloads/cifar/cifar_jax/workload.py

@@ -3,17 +3,16 @@
 import functools
 from typing import Any, Dict, Iterator, Optional, Tuple
 
-from flax import jax_utils
-from flax import linen as nn
-from flax.core import pop
 import jax
-from jax import lax
 import jax.numpy as jnp
 import optax
 import tensorflow_datasets as tfds
+from flax import jax_utils
+from flax import linen as nn
+from flax.core import pop
+from jax import lax
 
-from algoperf import param_utils
-from algoperf import spec
+from algoperf import param_utils, spec
 from algoperf.workloads.cifar.cifar_jax import models
 from algoperf.workloads.cifar.cifar_jax.input_pipeline import create_input_iter
 from algoperf.workloads.cifar.workload import BaseCifarWorkload