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,14 +1,13 @@
 import os
 from typing import Tuple
 
-from absl import logging
 import jax
 import tensorflow as tf
 import torch
-from torch import nn
-from torch import Tensor
 import torch.distributed as dist
 import torch.nn.functional as F
+from absl import logging
+from torch import Tensor, nn
 
 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/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_pytorch/models.py

@@ -14,11 +14,9 @@
 from algoperf.init_utils import pytorch_default_init
 from algoperf.workloads.imagenet_resnet.imagenet_pytorch.models import (
   BasicBlock,
-)
-from algoperf.workloads.imagenet_resnet.imagenet_pytorch.models import (
   Bottleneck,
+  conv1x1,
 )
-from algoperf.workloads.imagenet_resnet.imagenet_pytorch.models import conv1x1
 
 
 class ResNet(nn.Module):

algoperf/workloads/cifar/cifar_pytorch/workload.py

@@ -12,10 +12,7 @@
 from torchvision import transforms
 from torchvision.datasets import CIFAR10
 
-from algoperf import data_utils
-from algoperf import param_utils
-from algoperf import pytorch_utils
-from algoperf import spec
+from algoperf import data_utils, param_utils, pytorch_utils, spec
 from algoperf.workloads.cifar.cifar_pytorch.models import resnet18
 from algoperf.workloads.cifar.workload import BaseCifarWorkload