Add ufmt (usort + black) as code formatter (#4384)

* add ufmt as code formatter

* cleanup

* quote ufmt requirement

* split imports into more groups

* regenerate circleci config

* fix CI

* clarify local testing utils section

* use ufmt pre-commit hook

* split relative imports into local category

* Revert "split relative imports into local category"

This reverts commit f2e224cde2008c56c9347c1f69746d39065cdd51.

* pin black and usort dependencies

* fix local test utils detection

* fix ufmt rev

* add reference utils to local category

* fix usort config

* remove custom categories sorting

* Run pre-commit without fixing flake8

* got a double import in merge
Co-authored-by: Nicolas Hug <nicolashug@fb.com>

references/detection/group_by_aspect_ratio.py

 import bisect
-from collections import defaultdict
 import copy
-from itertools import repeat, chain
 import math
-import numpy as np
+from collections import defaultdict
+from itertools import repeat, chain
 
+import numpy as np
 import torch
 import torch.utils.data
-from torch.utils.data.sampler import BatchSampler, Sampler
-from torch.utils.model_zoo import tqdm
 import torchvision
-
 from PIL import Image
+from torch.utils.data.sampler import BatchSampler, Sampler
+from torch.utils.model_zoo import tqdm
 
 
 def _repeat_to_at_least(iterable, n):
@@ -34,11 +33,11 @@ class GroupedBatchSampler(BatchSampler):
             0, i.e. they must be in the range [0, num_groups).
         batch_size (int): Size of mini-batch.
     """
+
     def __init__(self, sampler, group_ids, batch_size):
         if not isinstance(sampler, Sampler):
             raise ValueError(
-                "sampler should be an instance of "
-                "torch.utils.data.Sampler, but got sampler={}".format(sampler)
+                "sampler should be an instance of " "torch.utils.data.Sampler, but got sampler={}".format(sampler)
             )
         self.sampler = sampler
         self.group_ids = group_ids
@@ -68,8 +67,7 @@ class GroupedBatchSampler(BatchSampler):
         if num_remaining > 0:
             # for the remaining batches, take first the buffers with largest number
             # of elements
-            for group_id, _ in sorted(buffer_per_group.items(),
-                                      key=lambda x: len(x[1]), reverse=True):
+            for group_id, _ in sorted(buffer_per_group.items(), key=lambda x: len(x[1]), reverse=True):
                 remaining = self.batch_size - len(buffer_per_group[group_id])
                 samples_from_group_id = _repeat_to_at_least(samples_per_group[group_id], remaining)
                 buffer_per_group[group_id].extend(samples_from_group_id[:remaining])
@@ -85,10 +83,12 @@ class GroupedBatchSampler(BatchSampler):
 
 
 def _compute_aspect_ratios_slow(dataset, indices=None):
-    print("Your dataset doesn't support the fast path for "
-          "computing the aspect ratios, so will iterate over "
-          "the full dataset and load every image instead. "
-          "This might take some time...")
+    print(
+        "Your dataset doesn't support the fast path for "
+        "computing the aspect ratios, so will iterate over "
+        "the full dataset and load every image instead. "
+        "This might take some time..."
+    )
     if indices is None:
         indices = range(len(dataset))
 
@@ -104,9 +104,12 @@ def _compute_aspect_ratios_slow(dataset, indices=None):
 
     sampler = SubsetSampler(indices)
     data_loader = torch.utils.data.DataLoader(
-        dataset, batch_size=1, sampler=sampler,
+        dataset,
+        batch_size=1,
+        sampler=sampler,
         num_workers=14,  # you might want to increase it for faster processing
-        collate_fn=lambda x: x[0])
+        collate_fn=lambda x: x[0],
+    )
     aspect_ratios = []
     with tqdm(total=len(dataset)) as pbar:
         for _i, (img, _) in enumerate(data_loader):

references/detection/presets.py

 import torch
-
 import transforms as T
 
 
 class DetectionPresetTrain:
-    def __init__(self, data_augmentation, hflip_prob=0.5, mean=(123., 117., 104.)):
-        if data_augmentation == 'hflip':
-            self.transforms = T.Compose([
-                T.RandomHorizontalFlip(p=hflip_prob),
-                T.PILToTensor(),
-                T.ConvertImageDtype(torch.float),
-            ])
-        elif data_augmentation == 'ssd':
-            self.transforms = T.Compose([
-                T.RandomPhotometricDistort(),
-                T.RandomZoomOut(fill=list(mean)),
-                T.RandomIoUCrop(),
-                T.RandomHorizontalFlip(p=hflip_prob),
-                T.PILToTensor(),
-                T.ConvertImageDtype(torch.float),
-            ])
-        elif data_augmentation == 'ssdlite':
-            self.transforms = T.Compose([
-                T.RandomIoUCrop(),
-                T.RandomHorizontalFlip(p=hflip_prob),
-                T.PILToTensor(),
-                T.ConvertImageDtype(torch.float),
-            ])
+    def __init__(self, data_augmentation, hflip_prob=0.5, mean=(123.0, 117.0, 104.0)):
+        if data_augmentation == "hflip":
+            self.transforms = T.Compose(
+                [
+                    T.RandomHorizontalFlip(p=hflip_prob),
+                    T.PILToTensor(),
+                    T.ConvertImageDtype(torch.float),
+                ]
+            )
+        elif data_augmentation == "ssd":
+            self.transforms = T.Compose(
+                [
+                    T.RandomPhotometricDistort(),
+                    T.RandomZoomOut(fill=list(mean)),
+                    T.RandomIoUCrop(),
+                    T.RandomHorizontalFlip(p=hflip_prob),
+                    T.PILToTensor(),
+                    T.ConvertImageDtype(torch.float),
+                ]
+            )
+        elif data_augmentation == "ssdlite":
+            self.transforms = T.Compose(
+                [
+                    T.RandomIoUCrop(),
+                    T.RandomHorizontalFlip(p=hflip_prob),
+                    T.PILToTensor(),
+                    T.ConvertImageDtype(torch.float),
+                ]
+            )
         else:
             raise ValueError(f'Unknown data augmentation policy "{data_augmentation}"')
 

references/detection/train.py

@@ -21,26 +21,20 @@ import datetime
 import os
 import time
 
+import presets
 import torch
 import torch.utils.data
 import torchvision
 import torchvision.models.detection
 import torchvision.models.detection.mask_rcnn
-
+import utils
 from coco_utils import get_coco, get_coco_kp
-
-from group_by_aspect_ratio import GroupedBatchSampler, create_aspect_ratio_groups
 from engine import train_one_epoch, evaluate
-
-import presets
-import utils
+from group_by_aspect_ratio import GroupedBatchSampler, create_aspect_ratio_groups
 
 
 def get_dataset(name, image_set, transform, data_path):
-    paths = {
-        "coco": (data_path, get_coco, 91),
-        "coco_kp": (data_path, get_coco_kp, 2)
-    }
+    paths = {"coco": (data_path, get_coco, 91), "coco_kp": (data_path, get_coco_kp, 2)}
     p, ds_fn, num_classes = paths[name]
 
     ds = ds_fn(p, image_set=image_set, transforms=transform)
@@ -53,42 +47,60 @@ def get_transform(train, data_augmentation):
 
 def get_args_parser(add_help=True):
     import argparse
-    parser = argparse.ArgumentParser(description='PyTorch Detection Training', add_help=add_help)
-
-    parser.add_argument('--data-path', default='/datasets01/COCO/022719/', help='dataset')
-    parser.add_argument('--dataset', default='coco', help='dataset')
-    parser.add_argument('--model', default='maskrcnn_resnet50_fpn', help='model')
-    parser.add_argument('--device', default='cuda', help='device')
-    parser.add_argument('-b', '--batch-size', default=2, type=int,
-                        help='images per gpu, the total batch size is $NGPU x batch_size')
-    parser.add_argument('--epochs', default=26, type=int, metavar='N',
-                        help='number of total epochs to run')
-    parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
-                        help='number of data loading workers (default: 4)')
-    parser.add_argument('--lr', default=0.02, type=float,
-                        help='initial learning rate, 0.02 is the default value for training '
-                             'on 8 gpus and 2 images_per_gpu')
-    parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
-                        help='momentum')
-    parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
-                        metavar='W', help='weight decay (default: 1e-4)',
-                        dest='weight_decay')
-    parser.add_argument('--lr-scheduler', default="multisteplr", help='the lr scheduler (default: multisteplr)')
-    parser.add_argument('--lr-step-size', default=8, type=int,
-                        help='decrease lr every step-size epochs (multisteplr scheduler only)')
-    parser.add_argument('--lr-steps', default=[16, 22], nargs='+', type=int,
-                        help='decrease lr every step-size epochs (multisteplr scheduler only)')
-    parser.add_argument('--lr-gamma', default=0.1, type=float,
-                        help='decrease lr by a factor of lr-gamma (multisteplr scheduler only)')
-    parser.add_argument('--print-freq', default=20, type=int, help='print frequency')
-    parser.add_argument('--output-dir', default='.', help='path where to save')
-    parser.add_argument('--resume', default='', help='resume from checkpoint')
-    parser.add_argument('--start_epoch', default=0, type=int, help='start epoch')
-    parser.add_argument('--aspect-ratio-group-factor', default=3, type=int)
-    parser.add_argument('--rpn-score-thresh', default=None, type=float, help='rpn score threshold for faster-rcnn')
-    parser.add_argument('--trainable-backbone-layers', default=None, type=int,
-                        help='number of trainable layers of backbone')
-    parser.add_argument('--data-augmentation', default="hflip", help='data augmentation policy (default: hflip)')
+
+    parser = argparse.ArgumentParser(description="PyTorch Detection Training", add_help=add_help)
+
+    parser.add_argument("--data-path", default="/datasets01/COCO/022719/", help="dataset")
+    parser.add_argument("--dataset", default="coco", help="dataset")
+    parser.add_argument("--model", default="maskrcnn_resnet50_fpn", help="model")
+    parser.add_argument("--device", default="cuda", help="device")
+    parser.add_argument(
+        "-b", "--batch-size", default=2, type=int, help="images per gpu, the total batch size is $NGPU x batch_size"
+    )
+    parser.add_argument("--epochs", default=26, type=int, metavar="N", help="number of total epochs to run")
+    parser.add_argument(
+        "-j", "--workers", default=4, type=int, metavar="N", help="number of data loading workers (default: 4)"
+    )
+    parser.add_argument(
+        "--lr",
+        default=0.02,
+        type=float,
+        help="initial learning rate, 0.02 is the default value for training " "on 8 gpus and 2 images_per_gpu",
+    )
+    parser.add_argument("--momentum", default=0.9, type=float, metavar="M", help="momentum")
+    parser.add_argument(
+        "--wd",
+        "--weight-decay",
+        default=1e-4,
+        type=float,
+        metavar="W",
+        help="weight decay (default: 1e-4)",
+        dest="weight_decay",
+    )
+    parser.add_argument("--lr-scheduler", default="multisteplr", help="the lr scheduler (default: multisteplr)")
+    parser.add_argument(
+        "--lr-step-size", default=8, type=int, help="decrease lr every step-size epochs (multisteplr scheduler only)"
+    )
+    parser.add_argument(
+        "--lr-steps",
+        default=[16, 22],
+        nargs="+",
+        type=int,
+        help="decrease lr every step-size epochs (multisteplr scheduler only)",
+    )
+    parser.add_argument(
+        "--lr-gamma", default=0.1, type=float, help="decrease lr by a factor of lr-gamma (multisteplr scheduler only)"
+    )
+    parser.add_argument("--print-freq", default=20, type=int, help="print frequency")
+    parser.add_argument("--output-dir", default=".", help="path where to save")
+    parser.add_argument("--resume", default="", help="resume from checkpoint")
+    parser.add_argument("--start_epoch", default=0, type=int, help="start epoch")
+    parser.add_argument("--aspect-ratio-group-factor", default=3, type=int)
+    parser.add_argument("--rpn-score-thresh", default=None, type=float, help="rpn score threshold for faster-rcnn")
+    parser.add_argument(
+        "--trainable-backbone-layers", default=None, type=int, help="number of trainable layers of backbone"
+    )
+    parser.add_argument("--data-augmentation", default="hflip", help="data augmentation policy (default: hflip)")
     parser.add_argument(
         "--sync-bn",
         dest="sync_bn",
@@ -109,9 +121,8 @@ def get_args_parser(add_help=True):
     )
 
     # distributed training parameters
-    parser.add_argument('--world-size', default=1, type=int,
-                        help='number of distributed processes')
-    parser.add_argument('--dist-url', default='env://', help='url used to set up distributed training')
+    parser.add_argument("--world-size", default=1, type=int, help="number of distributed processes")
+    parser.add_argument("--dist-url", default="env://", help="url used to set up distributed training")
 
     return parser
 
@@ -128,8 +139,9 @@ def main(args):
     # Data loading code
     print("Loading data")
 
-    dataset, num_classes = get_dataset(args.dataset, "train", get_transform(True, args.data_augmentation),
-                                       args.data_path)
+    dataset, num_classes = get_dataset(
+        args.dataset, "train", get_transform(True, args.data_augmentation), args.data_path
+    )
     dataset_test, _ = get_dataset(args.dataset, "val", get_transform(False, args.data_augmentation), args.data_path)
 
     print("Creating data loaders")
@@ -144,27 +156,24 @@ def main(args):
         group_ids = create_aspect_ratio_groups(dataset, k=args.aspect_ratio_group_factor)
         train_batch_sampler = GroupedBatchSampler(train_sampler, group_ids, args.batch_size)
     else:
-        train_batch_sampler = torch.utils.data.BatchSampler(
-            train_sampler, args.batch_size, drop_last=True)
+        train_batch_sampler = torch.utils.data.BatchSampler(train_sampler, args.batch_size, drop_last=True)
 
     data_loader = torch.utils.data.DataLoader(
-        dataset, batch_sampler=train_batch_sampler, num_workers=args.workers,
-        collate_fn=utils.collate_fn)
+        dataset, batch_sampler=train_batch_sampler, num_workers=args.workers, collate_fn=utils.collate_fn
+    )
 
     data_loader_test = torch.utils.data.DataLoader(
-        dataset_test, batch_size=1,
-        sampler=test_sampler, num_workers=args.workers,
-        collate_fn=utils.collate_fn)
+        dataset_test, batch_size=1, sampler=test_sampler, num_workers=args.workers, collate_fn=utils.collate_fn
+    )
 
     print("Creating model")
-    kwargs = {
-        "trainable_backbone_layers": args.trainable_backbone_layers
-    }
+    kwargs = {"trainable_backbone_layers": args.trainable_backbone_layers}
     if "rcnn" in args.model:
         if args.rpn_score_thresh is not None:
             kwargs["rpn_score_thresh"] = args.rpn_score_thresh
-    model = torchvision.models.detection.__dict__[args.model](num_classes=num_classes, pretrained=args.pretrained,
-                                                              **kwargs)
+    model = torchvision.models.detection.__dict__[args.model](
+        num_classes=num_classes, pretrained=args.pretrained, **kwargs
+    )
     model.to(device)
     if args.distributed and args.sync_bn:
         model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
@@ -175,24 +184,25 @@ def main(args):
         model_without_ddp = model.module
 
     params = [p for p in model.parameters() if p.requires_grad]
-    optimizer = torch.optim.SGD(
-        params, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
+    optimizer = torch.optim.SGD(params, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
 
     args.lr_scheduler = args.lr_scheduler.lower()
-    if args.lr_scheduler == 'multisteplr':
+    if args.lr_scheduler == "multisteplr":
         lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.lr_steps, gamma=args.lr_gamma)
-    elif args.lr_scheduler == 'cosineannealinglr':
+    elif args.lr_scheduler == "cosineannealinglr":
         lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)
     else:
-        raise RuntimeError("Invalid lr scheduler '{}'. Only MultiStepLR and CosineAnnealingLR "
-                           "are supported.".format(args.lr_scheduler))
+        raise RuntimeError(
+            "Invalid lr scheduler '{}'. Only MultiStepLR and CosineAnnealingLR "
+            "are supported.".format(args.lr_scheduler)
+        )
 
     if args.resume:
-        checkpoint = torch.load(args.resume, map_location='cpu')
-        model_without_ddp.load_state_dict(checkpoint['model'])
-        optimizer.load_state_dict(checkpoint['optimizer'])
-        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
-        args.start_epoch = checkpoint['epoch'] + 1
+        checkpoint = torch.load(args.resume, map_location="cpu")
+        model_without_ddp.load_state_dict(checkpoint["model"])
+        optimizer.load_state_dict(checkpoint["optimizer"])
+        lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
+        args.start_epoch = checkpoint["epoch"] + 1
 
     if args.test_only:
         evaluate(model, data_loader_test, device=device)
@@ -207,25 +217,21 @@ def main(args):
         lr_scheduler.step()
         if args.output_dir:
             checkpoint = {
-                'model': model_without_ddp.state_dict(),
-                'optimizer': optimizer.state_dict(),
-                'lr_scheduler': lr_scheduler.state_dict(),
-                'args': args,
-                'epoch': epoch
+                "model": model_without_ddp.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "lr_scheduler": lr_scheduler.state_dict(),
+                "args": args,
+                "epoch": epoch,
             }
-            utils.save_on_master(
-                checkpoint,
-                os.path.join(args.output_dir, 'model_{}.pth'.format(epoch)))
-            utils.save_on_master(
-                checkpoint,
-                os.path.join(args.output_dir, 'checkpoint.pth'))
+            utils.save_on_master(checkpoint, os.path.join(args.output_dir, "model_{}.pth".format(epoch)))
+            utils.save_on_master(checkpoint, os.path.join(args.output_dir, "checkpoint.pth"))
 
         # evaluate after every epoch
         evaluate(model, data_loader_test, device=device)
 
     total_time = time.time() - start_time
     total_time_str = str(datetime.timedelta(seconds=int(total_time)))
-    print('Training time {}'.format(total_time_str))
+    print("Training time {}".format(total_time_str))
 
 
 if __name__ == "__main__":

references/detection/transforms.py

@@ -28,8 +28,9 @@ class Compose(object):
 
 
 class RandomHorizontalFlip(T.RandomHorizontalFlip):
-    def forward(self, image: Tensor,
-                target: Optional[Dict[str, Tensor]] = None) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
         if torch.rand(1) < self.p:
             image = F.hflip(image)
             if target is not None:
@@ -45,16 +46,18 @@ class RandomHorizontalFlip(T.RandomHorizontalFlip):
 
 
 class ToTensor(nn.Module):
-    def forward(self, image: Tensor,
-                target: Optional[Dict[str, Tensor]] = None) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
         image = F.pil_to_tensor(image)
         image = F.convert_image_dtype(image)
         return image, target
 
 
 class PILToTensor(nn.Module):
-    def forward(self, image: Tensor,
-                target: Optional[Dict[str, Tensor]] = None) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
         image = F.pil_to_tensor(image)
         return image, target
 
@@ -64,15 +67,23 @@ class ConvertImageDtype(nn.Module):
         super().__init__()
         self.dtype = dtype
 
-    def forward(self, image: Tensor,
-                target: Optional[Dict[str, Tensor]] = None) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
         image = F.convert_image_dtype(image, self.dtype)
         return image, target
 
 
 class RandomIoUCrop(nn.Module):
-    def __init__(self, min_scale: float = 0.3, max_scale: float = 1.0, min_aspect_ratio: float = 0.5,
-                 max_aspect_ratio: float = 2.0, sampler_options: Optional[List[float]] = None, trials: int = 40):
+    def __init__(
+        self,
+        min_scale: float = 0.3,
+        max_scale: float = 1.0,
+        min_aspect_ratio: float = 0.5,
+        max_aspect_ratio: float = 2.0,
+        sampler_options: Optional[List[float]] = None,
+        trials: int = 40,
+    ):
         super().__init__()
         # Configuration similar to https://github.com/weiliu89/caffe/blob/ssd/examples/ssd/ssd_coco.py#L89-L174
         self.min_scale = min_scale
@@ -84,14 +95,15 @@ class RandomIoUCrop(nn.Module):
         self.options = sampler_options
         self.trials = trials
 
-    def forward(self, image: Tensor,
-                target: Optional[Dict[str, Tensor]] = None) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
         if target is None:
             raise ValueError("The targets can't be None for this transform.")
 
         if isinstance(image, torch.Tensor):
             if image.ndimension() not in {2, 3}:
-                raise ValueError('image should be 2/3 dimensional. Got {} dimensions.'.format(image.ndimension()))
+                raise ValueError("image should be 2/3 dimensional. Got {} dimensions.".format(image.ndimension()))
             elif image.ndimension() == 2:
                 image = image.unsqueeze(0)
 
@@ -131,8 +143,9 @@ class RandomIoUCrop(nn.Module):
 
                 # check at least 1 box with jaccard limitations
                 boxes = target["boxes"][is_within_crop_area]
-                ious = torchvision.ops.boxes.box_iou(boxes, torch.tensor([[left, top, right, bottom]],
-                                                                         dtype=boxes.dtype, device=boxes.device))
+                ious = torchvision.ops.boxes.box_iou(
+                    boxes, torch.tensor([[left, top, right, bottom]], dtype=boxes.dtype, device=boxes.device)
+                )
                 if ious.max() < min_jaccard_overlap:
                     continue
 
@@ -149,13 +162,15 @@ class RandomIoUCrop(nn.Module):
 
 
 class RandomZoomOut(nn.Module):
-    def __init__(self, fill: Optional[List[float]] = None, side_range: Tuple[float, float] = (1., 4.), p: float = 0.5):
+    def __init__(
+        self, fill: Optional[List[float]] = None, side_range: Tuple[float, float] = (1.0, 4.0), p: float = 0.5
+    ):
         super().__init__()
         if fill is None:
-            fill = [0., 0., 0.]
+            fill = [0.0, 0.0, 0.0]
         self.fill = fill
         self.side_range = side_range
-        if side_range[0] < 1. or side_range[0] > side_range[1]:
+        if side_range[0] < 1.0 or side_range[0] > side_range[1]:
             raise ValueError("Invalid canvas side range provided {}.".format(side_range))
         self.p = p
 
@@ -165,11 +180,12 @@ class RandomZoomOut(nn.Module):
         # We fake the type to make it work on JIT
         return tuple(int(x) for x in self.fill) if is_pil else 0
 
-    def forward(self, image: Tensor,
-                target: Optional[Dict[str, Tensor]] = None) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
         if isinstance(image, torch.Tensor):
             if image.ndimension() not in {2, 3}:
-                raise ValueError('image should be 2/3 dimensional. Got {} dimensions.'.format(image.ndimension()))
+                raise ValueError("image should be 2/3 dimensional. Got {} dimensions.".format(image.ndimension()))
             elif image.ndimension() == 2:
                 image = image.unsqueeze(0)
 
@@ -196,8 +212,9 @@ class RandomZoomOut(nn.Module):
         image = F.pad(image, [left, top, right, bottom], fill=fill)
         if isinstance(image, torch.Tensor):
             v = torch.tensor(self.fill, device=image.device, dtype=image.dtype).view(-1, 1, 1)
-            image[..., :top, :] = image[..., :, :left] = image[..., (top + orig_h):, :] = \
-                image[..., :, (left + orig_w):] = v
+            image[..., :top, :] = image[..., :, :left] = image[..., (top + orig_h) :, :] = image[
+                ..., :, (left + orig_w) :
+            ] = v
 
         if target is not None:
             target["boxes"][:, 0::2] += left
@@ -207,8 +224,14 @@ class RandomZoomOut(nn.Module):
 
 
 class RandomPhotometricDistort(nn.Module):
-    def __init__(self, contrast: Tuple[float] = (0.5, 1.5), saturation: Tuple[float] = (0.5, 1.5),
-                 hue: Tuple[float] = (-0.05, 0.05), brightness: Tuple[float] = (0.875, 1.125), p: float = 0.5):
+    def __init__(
+        self,
+        contrast: Tuple[float] = (0.5, 1.5),
+        saturation: Tuple[float] = (0.5, 1.5),
+        hue: Tuple[float] = (-0.05, 0.05),
+        brightness: Tuple[float] = (0.875, 1.125),
+        p: float = 0.5,
+    ):
         super().__init__()
         self._brightness = T.ColorJitter(brightness=brightness)
         self._contrast = T.ColorJitter(contrast=contrast)
@@ -216,11 +239,12 @@ class RandomPhotometricDistort(nn.Module):
         self._saturation = T.ColorJitter(saturation=saturation)
         self.p = p
 
-    def forward(self, image: Tensor,
-                target: Optional[Dict[str, Tensor]] = None) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
         if isinstance(image, torch.Tensor):
             if image.ndimension() not in {2, 3}:
-                raise ValueError('image should be 2/3 dimensional. Got {} dimensions.'.format(image.ndimension()))
+                raise ValueError("image should be 2/3 dimensional. Got {} dimensions.".format(image.ndimension()))
             elif image.ndimension() == 2:
                 image = image.unsqueeze(0)
 

references/detection/utils.py

-from collections import defaultdict, deque
 import datetime
 import errno
 import os
 import time
+from collections import defaultdict, deque
 
 import torch
 import torch.distributed as dist
@@ -32,7 +32,7 @@ class SmoothedValue(object):
         """
         if not is_dist_avail_and_initialized():
             return
-        t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda")
         dist.barrier()
         dist.all_reduce(t)
         t = t.tolist()
@@ -63,11 +63,8 @@ class SmoothedValue(object):
 
     def __str__(self):
         return self.fmt.format(
-            median=self.median,
-            avg=self.avg,
-            global_avg=self.global_avg,
-            max=self.max,
-            value=self.value)
+            median=self.median, avg=self.avg, global_avg=self.global_avg, max=self.max, value=self.value
+        )
 
 
 def all_gather(data):
@@ -130,15 +127,12 @@ class MetricLogger(object):
             return self.meters[attr]
         if attr in self.__dict__:
             return self.__dict__[attr]
-        raise AttributeError("'{}' object has no attribute '{}'".format(
-            type(self).__name__, attr))
+        raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, attr))
 
     def __str__(self):
         loss_str = []
         for name, meter in self.meters.items():
-            loss_str.append(
-                "{}: {}".format(name, str(meter))
-            )
+            loss_str.append("{}: {}".format(name, str(meter)))
         return self.delimiter.join(loss_str)
 
     def synchronize_between_processes(self):
@@ -151,31 +145,28 @@ class MetricLogger(object):
     def log_every(self, iterable, print_freq, header=None):
         i = 0
         if not header:
-            header = ''
+            header = ""
         start_time = time.time()
         end = time.time()
-        iter_time = SmoothedValue(fmt='{avg:.4f}')
-        data_time = SmoothedValue(fmt='{avg:.4f}')
-        space_fmt = ':' + str(len(str(len(iterable)))) + 'd'
+        iter_time = SmoothedValue(fmt="{avg:.4f}")
+        data_time = SmoothedValue(fmt="{avg:.4f}")
+        space_fmt = ":" + str(len(str(len(iterable)))) + "d"
         if torch.cuda.is_available():
-            log_msg = self.delimiter.join([
-                header,
-                '[{0' + space_fmt + '}/{1}]',
-                'eta: {eta}',
-                '{meters}',
-                'time: {time}',
-                'data: {data}',
-                'max mem: {memory:.0f}'
-            ])
+            log_msg = self.delimiter.join(
+                [
+                    header,
+                    "[{0" + space_fmt + "}/{1}]",
+                    "eta: {eta}",
+                    "{meters}",
+                    "time: {time}",
+                    "data: {data}",
+                    "max mem: {memory:.0f}",
+                ]
+            )
         else:
-            log_msg = self.delimiter.join([
-                header,
-                '[{0' + space_fmt + '}/{1}]',
-                'eta: {eta}',
-                '{meters}',
-                'time: {time}',
-                'data: {data}'
-            ])
+            log_msg = self.delimiter.join(
+                [header, "[{0" + space_fmt + "}/{1}]", "eta: {eta}", "{meters}", "time: {time}", "data: {data}"]
+            )
         MB = 1024.0 * 1024.0
         for obj in iterable:
             data_time.update(time.time() - end)
@@ -185,22 +176,28 @@ class MetricLogger(object):
                 eta_seconds = iter_time.global_avg * (len(iterable) - i)
                 eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
                 if torch.cuda.is_available():
-                    print(log_msg.format(
-                        i, len(iterable), eta=eta_string,
-                        meters=str(self),
-                        time=str(iter_time), data=str(data_time),
-                        memory=torch.cuda.max_memory_allocated() / MB))
+                    print(
+                        log_msg.format(
+                            i,
+                            len(iterable),
+                            eta=eta_string,
+                            meters=str(self),
+                            time=str(iter_time),
+                            data=str(data_time),
+                            memory=torch.cuda.max_memory_allocated() / MB,
+                        )
+                    )
                 else:
-                    print(log_msg.format(
-                        i, len(iterable), eta=eta_string,
-                        meters=str(self),
-                        time=str(iter_time), data=str(data_time)))
+                    print(
+                        log_msg.format(
+                            i, len(iterable), eta=eta_string, meters=str(self), time=str(iter_time), data=str(data_time)
+                        )
+                    )
             i += 1
             end = time.time()
         total_time = time.time() - start_time
         total_time_str = str(datetime.timedelta(seconds=int(total_time)))
-        print('{} Total time: {} ({:.4f} s / it)'.format(
-            header, total_time_str, total_time / len(iterable)))
+        print("{} Total time: {} ({:.4f} s / it)".format(header, total_time_str, total_time / len(iterable)))
 
 
 def collate_fn(batch):
@@ -220,10 +217,11 @@ def setup_for_distributed(is_master):
     This function disables printing when not in master process
     """
     import builtins as __builtin__
+
     builtin_print = __builtin__.print
 
     def print(*args, **kwargs):
-        force = kwargs.pop('force', False)
+        force = kwargs.pop("force", False)
         if is_master or force:
             builtin_print(*args, **kwargs)
 
@@ -260,25 +258,25 @@ def save_on_master(*args, **kwargs):
 
 
 def init_distributed_mode(args):
-    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
+    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
         args.rank = int(os.environ["RANK"])
-        args.world_size = int(os.environ['WORLD_SIZE'])
-        args.gpu = int(os.environ['LOCAL_RANK'])
-    elif 'SLURM_PROCID' in os.environ:
-        args.rank = int(os.environ['SLURM_PROCID'])
+        args.world_size = int(os.environ["WORLD_SIZE"])
+        args.gpu = int(os.environ["LOCAL_RANK"])
+    elif "SLURM_PROCID" in os.environ:
+        args.rank = int(os.environ["SLURM_PROCID"])
         args.gpu = args.rank % torch.cuda.device_count()
     else:
-        print('Not using distributed mode')
+        print("Not using distributed mode")
         args.distributed = False
         return
 
     args.distributed = True
 
     torch.cuda.set_device(args.gpu)
-    args.dist_backend = 'nccl'
-    print('| distributed init (rank {}): {}'.format(
-        args.rank, args.dist_url), flush=True)
-    torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
-                                         world_size=args.world_size, rank=args.rank)
+    args.dist_backend = "nccl"
+    print("| distributed init (rank {}): {}".format(args.rank, args.dist_url), flush=True)
+    torch.distributed.init_process_group(
+        backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank
+    )
     torch.distributed.barrier()
     setup_for_distributed(args.rank == 0)

references/segmentation/coco_utils.py

 import copy
+import os
+
 import torch
 import torch.utils.data
 import torchvision
 from PIL import Image
-
-import os
-
 from pycocotools import mask as coco_mask
-
 from transforms import Compose
 
 
@@ -90,14 +88,9 @@ def get_coco(root, image_set, transforms):
         "val": ("val2017", os.path.join("annotations", "instances_val2017.json")),
         # "train": ("val2017", os.path.join("annotations", "instances_val2017.json"))
     }
-    CAT_LIST = [0, 5, 2, 16, 9, 44, 6, 3, 17, 62, 21, 67, 18, 19, 4,
-                1, 64, 20, 63, 7, 72]
-
-    transforms = Compose([
-        FilterAndRemapCocoCategories(CAT_LIST, remap=True),
-        ConvertCocoPolysToMask(),
-        transforms
-    ])
+    CAT_LIST = [0, 5, 2, 16, 9, 44, 6, 3, 17, 62, 21, 67, 18, 19, 4, 1, 64, 20, 63, 7, 72]
+
+    transforms = Compose([FilterAndRemapCocoCategories(CAT_LIST, remap=True), ConvertCocoPolysToMask(), transforms])
 
     img_folder, ann_file = PATHS[image_set]
     img_folder = os.path.join(root, img_folder)

references/segmentation/presets.py

 import torch
-
 import transforms as T
 
 
@@ -11,12 +10,14 @@ class SegmentationPresetTrain:
         trans = [T.RandomResize(min_size, max_size)]
         if hflip_prob > 0:
             trans.append(T.RandomHorizontalFlip(hflip_prob))
-        trans.extend([
-            T.RandomCrop(crop_size),
-            T.PILToTensor(),
-            T.ConvertImageDtype(torch.float),
-            T.Normalize(mean=mean, std=std),
-        ])
+        trans.extend(
+            [
+                T.RandomCrop(crop_size),
+                T.PILToTensor(),
+                T.ConvertImageDtype(torch.float),
+                T.Normalize(mean=mean, std=std),
+            ]
+        )
         self.transforms = T.Compose(trans)
 
     def __call__(self, img, target):
@@ -25,12 +26,14 @@ class SegmentationPresetTrain:
 
 class SegmentationPresetEval:
     def __init__(self, base_size, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
-        self.transforms = T.Compose([
-            T.RandomResize(base_size, base_size),
-            T.PILToTensor(),
-            T.ConvertImageDtype(torch.float),
-            T.Normalize(mean=mean, std=std),
-        ])
+        self.transforms = T.Compose(
+            [
+                T.RandomResize(base_size, base_size),
+                T.PILToTensor(),
+                T.ConvertImageDtype(torch.float),
+                T.Normalize(mean=mean, std=std),
+            ]
+        )
 
     def __call__(self, img, target):
         return self.transforms(img, target)

references/segmentation/train.py

@@ -2,23 +2,23 @@ import datetime
 import os
 import time
 
+import presets
 import torch
 import torch.utils.data
-from torch import nn
 import torchvision
-
-from coco_utils import get_coco
-import presets
 import utils
+from coco_utils import get_coco
+from torch import nn
 
 
 def get_dataset(dir_path, name, image_set, transform):
     def sbd(*args, **kwargs):
-        return torchvision.datasets.SBDataset(*args, mode='segmentation', **kwargs)
+        return torchvision.datasets.SBDataset(*args, mode="segmentation", **kwargs)
+
     paths = {
         "voc": (dir_path, torchvision.datasets.VOCSegmentation, 21),
         "voc_aug": (dir_path, sbd, 21),
-        "coco": (dir_path, get_coco, 21)
+        "coco": (dir_path, get_coco, 21),
     }
     p, ds_fn, num_classes = paths[name]
 
@@ -39,21 +39,21 @@ def criterion(inputs, target):
         losses[name] = nn.functional.cross_entropy(x, target, ignore_index=255)
 
     if len(losses) == 1:
-        return losses['out']
+        return losses["out"]
 
-    return losses['out'] + 0.5 * losses['aux']
+    return losses["out"] + 0.5 * losses["aux"]
 
 
 def evaluate(model, data_loader, device, num_classes):
     model.eval()
     confmat = utils.ConfusionMatrix(num_classes)
     metric_logger = utils.MetricLogger(delimiter="  ")
-    header = 'Test:'
+    header = "Test:"
     with torch.no_grad():
         for image, target in metric_logger.log_every(data_loader, 100, header):
             image, target = image.to(device), target.to(device)
             output = model(image)
-            output = output['out']
+            output = output["out"]
 
             confmat.update(target.flatten(), output.argmax(1).flatten())
 
@@ -65,8 +65,8 @@ def evaluate(model, data_loader, device, num_classes):
 def train_one_epoch(model, criterion, optimizer, data_loader, lr_scheduler, device, epoch, print_freq):
     model.train()
     metric_logger = utils.MetricLogger(delimiter="  ")
-    metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value}'))
-    header = 'Epoch: [{}]'.format(epoch)
+    metric_logger.add_meter("lr", utils.SmoothedValue(window_size=1, fmt="{value}"))
+    header = "Epoch: [{}]".format(epoch)
     for image, target in metric_logger.log_every(data_loader, print_freq, header):
         image, target = image.to(device), target.to(device)
         output = model(image)
@@ -101,18 +101,21 @@ def main(args):
         test_sampler = torch.utils.data.SequentialSampler(dataset_test)
 
     data_loader = torch.utils.data.DataLoader(
-        dataset, batch_size=args.batch_size,
-        sampler=train_sampler, num_workers=args.workers,
-        collate_fn=utils.collate_fn, drop_last=True)
+        dataset,
+        batch_size=args.batch_size,
+        sampler=train_sampler,
+        num_workers=args.workers,
+        collate_fn=utils.collate_fn,
+        drop_last=True,
+    )
 
     data_loader_test = torch.utils.data.DataLoader(
-        dataset_test, batch_size=1,
-        sampler=test_sampler, num_workers=args.workers,
-        collate_fn=utils.collate_fn)
+        dataset_test, batch_size=1, sampler=test_sampler, num_workers=args.workers, collate_fn=utils.collate_fn
+    )
 
-    model = torchvision.models.segmentation.__dict__[args.model](num_classes=num_classes,
-                                                                 aux_loss=args.aux_loss,
-                                                                 pretrained=args.pretrained)
+    model = torchvision.models.segmentation.__dict__[args.model](
+        num_classes=num_classes, aux_loss=args.aux_loss, pretrained=args.pretrained
+    )
     model.to(device)
     if args.distributed:
         model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
@@ -129,42 +132,42 @@ def main(args):
     if args.aux_loss:
         params = [p for p in model_without_ddp.aux_classifier.parameters() if p.requires_grad]
         params_to_optimize.append({"params": params, "lr": args.lr * 10})
-    optimizer = torch.optim.SGD(
-        params_to_optimize,
-        lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
+    optimizer = torch.optim.SGD(params_to_optimize, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
 
     iters_per_epoch = len(data_loader)
     main_lr_scheduler = torch.optim.lr_scheduler.LambdaLR(
-        optimizer,
-        lambda x: (1 - x / (iters_per_epoch * (args.epochs - args.lr_warmup_epochs))) ** 0.9)
+        optimizer, lambda x: (1 - x / (iters_per_epoch * (args.epochs - args.lr_warmup_epochs))) ** 0.9
+    )
 
     if args.lr_warmup_epochs > 0:
         warmup_iters = iters_per_epoch * args.lr_warmup_epochs
         args.lr_warmup_method = args.lr_warmup_method.lower()
-        if args.lr_warmup_method == 'linear':
-            warmup_lr_scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=args.lr_warmup_decay,
-                                                                    total_iters=warmup_iters)
-        elif args.lr_warmup_method == 'constant':
-            warmup_lr_scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer, factor=args.lr_warmup_decay,
-                                                                      total_iters=warmup_iters)
+        if args.lr_warmup_method == "linear":
+            warmup_lr_scheduler = torch.optim.lr_scheduler.LinearLR(
+                optimizer, start_factor=args.lr_warmup_decay, total_iters=warmup_iters
+            )
+        elif args.lr_warmup_method == "constant":
+            warmup_lr_scheduler = torch.optim.lr_scheduler.ConstantLR(
+                optimizer, factor=args.lr_warmup_decay, total_iters=warmup_iters
+            )
         else:
-            raise RuntimeError("Invalid warmup lr method '{}'. Only linear and constant "
-                               "are supported.".format(args.lr_warmup_method))
+            raise RuntimeError(
+                "Invalid warmup lr method '{}'. Only linear and constant "
+                "are supported.".format(args.lr_warmup_method)
+            )
         lr_scheduler = torch.optim.lr_scheduler.SequentialLR(
-            optimizer,
-            schedulers=[warmup_lr_scheduler, main_lr_scheduler],
-            milestones=[warmup_iters]
+            optimizer, schedulers=[warmup_lr_scheduler, main_lr_scheduler], milestones=[warmup_iters]
         )
     else:
         lr_scheduler = main_lr_scheduler
 
     if args.resume:
-        checkpoint = torch.load(args.resume, map_location='cpu')
-        model_without_ddp.load_state_dict(checkpoint['model'], strict=not args.test_only)
+        checkpoint = torch.load(args.resume, map_location="cpu")
+        model_without_ddp.load_state_dict(checkpoint["model"], strict=not args.test_only)
         if not args.test_only:
-            optimizer.load_state_dict(checkpoint['optimizer'])
-            lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
-            args.start_epoch = checkpoint['epoch'] + 1
+            optimizer.load_state_dict(checkpoint["optimizer"])
+            lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
+            args.start_epoch = checkpoint["epoch"] + 1
 
     if args.test_only:
         confmat = evaluate(model, data_loader_test, device=device, num_classes=num_classes)
@@ -179,53 +182,54 @@ def main(args):
         confmat = evaluate(model, data_loader_test, device=device, num_classes=num_classes)
         print(confmat)
         checkpoint = {
-            'model': model_without_ddp.state_dict(),
-            'optimizer': optimizer.state_dict(),
-            'lr_scheduler': lr_scheduler.state_dict(),
-            'epoch': epoch,
-            'args': args
+            "model": model_without_ddp.state_dict(),
+            "optimizer": optimizer.state_dict(),
+            "lr_scheduler": lr_scheduler.state_dict(),
+            "epoch": epoch,
+            "args": args,
         }
-        utils.save_on_master(
-            checkpoint,
-            os.path.join(args.output_dir, 'model_{}.pth'.format(epoch)))
-        utils.save_on_master(
-            checkpoint,
-            os.path.join(args.output_dir, 'checkpoint.pth'))
+        utils.save_on_master(checkpoint, os.path.join(args.output_dir, "model_{}.pth".format(epoch)))
+        utils.save_on_master(checkpoint, os.path.join(args.output_dir, "checkpoint.pth"))
 
     total_time = time.time() - start_time
     total_time_str = str(datetime.timedelta(seconds=int(total_time)))
-    print('Training time {}'.format(total_time_str))
+    print("Training time {}".format(total_time_str))
 
 
 def get_args_parser(add_help=True):
     import argparse
-    parser = argparse.ArgumentParser(description='PyTorch Segmentation Training', add_help=add_help)
-
-    parser.add_argument('--data-path', default='/datasets01/COCO/022719/', help='dataset path')
-    parser.add_argument('--dataset', default='coco', help='dataset name')
-    parser.add_argument('--model', default='fcn_resnet101', help='model')
-    parser.add_argument('--aux-loss', action='store_true', help='auxiliar loss')
-    parser.add_argument('--device', default='cuda', help='device')
-    parser.add_argument('-b', '--batch-size', default=8, type=int)
-    parser.add_argument('--epochs', default=30, type=int, metavar='N',
-                        help='number of total epochs to run')
-
-    parser.add_argument('-j', '--workers', default=16, type=int, metavar='N',
-                        help='number of data loading workers (default: 16)')
-    parser.add_argument('--lr', default=0.01, type=float, help='initial learning rate')
-    parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
-                        help='momentum')
-    parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
-                        metavar='W', help='weight decay (default: 1e-4)',
-                        dest='weight_decay')
-    parser.add_argument('--lr-warmup-epochs', default=0, type=int, help='the number of epochs to warmup (default: 0)')
-    parser.add_argument('--lr-warmup-method', default="linear", type=str, help='the warmup method (default: linear)')
-    parser.add_argument('--lr-warmup-decay', default=0.01, type=float, help='the decay for lr')
-    parser.add_argument('--print-freq', default=10, type=int, help='print frequency')
-    parser.add_argument('--output-dir', default='.', help='path where to save')
-    parser.add_argument('--resume', default='', help='resume from checkpoint')
-    parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
-                        help='start epoch')
+
+    parser = argparse.ArgumentParser(description="PyTorch Segmentation Training", add_help=add_help)
+
+    parser.add_argument("--data-path", default="/datasets01/COCO/022719/", help="dataset path")
+    parser.add_argument("--dataset", default="coco", help="dataset name")
+    parser.add_argument("--model", default="fcn_resnet101", help="model")
+    parser.add_argument("--aux-loss", action="store_true", help="auxiliar loss")
+    parser.add_argument("--device", default="cuda", help="device")
+    parser.add_argument("-b", "--batch-size", default=8, type=int)
+    parser.add_argument("--epochs", default=30, type=int, metavar="N", help="number of total epochs to run")
+
+    parser.add_argument(
+        "-j", "--workers", default=16, type=int, metavar="N", help="number of data loading workers (default: 16)"
+    )
+    parser.add_argument("--lr", default=0.01, type=float, help="initial learning rate")
+    parser.add_argument("--momentum", default=0.9, type=float, metavar="M", help="momentum")
+    parser.add_argument(
+        "--wd",
+        "--weight-decay",
+        default=1e-4,
+        type=float,
+        metavar="W",
+        help="weight decay (default: 1e-4)",
+        dest="weight_decay",
+    )
+    parser.add_argument("--lr-warmup-epochs", default=0, type=int, help="the number of epochs to warmup (default: 0)")
+    parser.add_argument("--lr-warmup-method", default="linear", type=str, help="the warmup method (default: linear)")
+    parser.add_argument("--lr-warmup-decay", default=0.01, type=float, help="the decay for lr")
+    parser.add_argument("--print-freq", default=10, type=int, help="print frequency")
+    parser.add_argument("--output-dir", default=".", help="path where to save")
+    parser.add_argument("--resume", default="", help="resume from checkpoint")
+    parser.add_argument("--start-epoch", default=0, type=int, metavar="N", help="start epoch")
     parser.add_argument(
         "--test-only",
         dest="test_only",
@@ -239,9 +243,8 @@ def get_args_parser(add_help=True):
         action="store_true",
     )
     # distributed training parameters
-    parser.add_argument('--world-size', default=1, type=int,
-                        help='number of distributed processes')
-    parser.add_argument('--dist-url', default='env://', help='url used to set up distributed training')
+    parser.add_argument("--world-size", default=1, type=int, help="number of distributed processes")
+    parser.add_argument("--dist-url", default="env://", help="url used to set up distributed training")
 
     return parser
 

references/segmentation/utils.py

-from collections import defaultdict, deque
 import datetime
+import errno
+import os
 import time
+from collections import defaultdict, deque
+
 import torch
 import torch.distributed as dist
 
-import errno
-import os
-
 
 class SmoothedValue(object):
     """Track a series of values and provide access to smoothed values over a
@@ -32,7 +32,7 @@ class SmoothedValue(object):
         """
         if not is_dist_avail_and_initialized():
             return
-        t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda")
         dist.barrier()
         dist.all_reduce(t)
         t = t.tolist()
@@ -63,11 +63,8 @@ class SmoothedValue(object):
 
     def __str__(self):
         return self.fmt.format(
-            median=self.median,
-            avg=self.avg,
-            global_avg=self.global_avg,
-            max=self.max,
-            value=self.value)
+            median=self.median, avg=self.avg, global_avg=self.global_avg, max=self.max, value=self.value
+        )
 
 
 class ConfusionMatrix(object):
@@ -82,7 +79,7 @@ class ConfusionMatrix(object):
         with torch.no_grad():
             k = (a >= 0) & (a < n)
             inds = n * a[k].to(torch.int64) + b[k]
-            self.mat += torch.bincount(inds, minlength=n**2).reshape(n, n)
+            self.mat += torch.bincount(inds, minlength=n ** 2).reshape(n, n)
 
     def reset(self):
         self.mat.zero_()
@@ -104,15 +101,12 @@ class ConfusionMatrix(object):
 
     def __str__(self):
         acc_global, acc, iu = self.compute()
-        return (
-            'global correct: {:.1f}\n'
-            'average row correct: {}\n'
-            'IoU: {}\n'
-            'mean IoU: {:.1f}').format(
-                acc_global.item() * 100,
-                ['{:.1f}'.format(i) for i in (acc * 100).tolist()],
-                ['{:.1f}'.format(i) for i in (iu * 100).tolist()],
-                iu.mean().item() * 100)
+        return ("global correct: {:.1f}\n" "average row correct: {}\n" "IoU: {}\n" "mean IoU: {:.1f}").format(
+            acc_global.item() * 100,
+            ["{:.1f}".format(i) for i in (acc * 100).tolist()],
+            ["{:.1f}".format(i) for i in (iu * 100).tolist()],
+            iu.mean().item() * 100,
+        )
 
 
 class MetricLogger(object):
@@ -132,15 +126,12 @@ class MetricLogger(object):
             return self.meters[attr]
         if attr in self.__dict__:
             return self.__dict__[attr]
-        raise AttributeError("'{}' object has no attribute '{}'".format(
-            type(self).__name__, attr))
+        raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, attr))
 
     def __str__(self):
         loss_str = []
         for name, meter in self.meters.items():
-            loss_str.append(
-                "{}: {}".format(name, str(meter))
-            )
+            loss_str.append("{}: {}".format(name, str(meter)))
         return self.delimiter.join(loss_str)
 
     def synchronize_between_processes(self):
@@ -153,31 +144,28 @@ class MetricLogger(object):
     def log_every(self, iterable, print_freq, header=None):
         i = 0
         if not header:
-            header = ''
+            header = ""
         start_time = time.time()
         end = time.time()
-        iter_time = SmoothedValue(fmt='{avg:.4f}')
-        data_time = SmoothedValue(fmt='{avg:.4f}')
-        space_fmt = ':' + str(len(str(len(iterable)))) + 'd'
+        iter_time = SmoothedValue(fmt="{avg:.4f}")
+        data_time = SmoothedValue(fmt="{avg:.4f}")
+        space_fmt = ":" + str(len(str(len(iterable)))) + "d"
         if torch.cuda.is_available():
-            log_msg = self.delimiter.join([
-                header,
-                '[{0' + space_fmt + '}/{1}]',
-                'eta: {eta}',
-                '{meters}',
-                'time: {time}',
-                'data: {data}',
-                'max mem: {memory:.0f}'
-            ])
+            log_msg = self.delimiter.join(
+                [
+                    header,
+                    "[{0" + space_fmt + "}/{1}]",
+                    "eta: {eta}",
+                    "{meters}",
+                    "time: {time}",
+                    "data: {data}",
+                    "max mem: {memory:.0f}",
+                ]
+            )
         else:
-            log_msg = self.delimiter.join([
-                header,
-                '[{0' + space_fmt + '}/{1}]',
-                'eta: {eta}',
-                '{meters}',
-                'time: {time}',
-                'data: {data}'
-            ])
+            log_msg = self.delimiter.join(
+                [header, "[{0" + space_fmt + "}/{1}]", "eta: {eta}", "{meters}", "time: {time}", "data: {data}"]
+            )
         MB = 1024.0 * 1024.0
         for obj in iterable:
             data_time.update(time.time() - end)
@@ -187,21 +175,28 @@ class MetricLogger(object):
                 eta_seconds = iter_time.global_avg * (len(iterable) - i)
                 eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
                 if torch.cuda.is_available():
-                    print(log_msg.format(
-                        i, len(iterable), eta=eta_string,
-                        meters=str(self),
-                        time=str(iter_time), data=str(data_time),
-                        memory=torch.cuda.max_memory_allocated() / MB))
+                    print(
+                        log_msg.format(
+                            i,
+                            len(iterable),
+                            eta=eta_string,
+                            meters=str(self),
+                            time=str(iter_time),
+                            data=str(data_time),
+                            memory=torch.cuda.max_memory_allocated() / MB,
+                        )
+                    )
                 else:
-                    print(log_msg.format(
-                        i, len(iterable), eta=eta_string,
-                        meters=str(self),
-                        time=str(iter_time), data=str(data_time)))
+                    print(
+                        log_msg.format(
+                            i, len(iterable), eta=eta_string, meters=str(self), time=str(iter_time), data=str(data_time)
+                        )
+                    )
             i += 1
             end = time.time()
         total_time = time.time() - start_time
         total_time_str = str(datetime.timedelta(seconds=int(total_time)))
-        print('{} Total time: {}'.format(header, total_time_str))
+        print("{} Total time: {}".format(header, total_time_str))
 
 
 def cat_list(images, fill_value=0):
@@ -209,7 +204,7 @@ def cat_list(images, fill_value=0):
     batch_shape = (len(images),) + max_size
     batched_imgs = images[0].new(*batch_shape).fill_(fill_value)
     for img, pad_img in zip(images, batched_imgs):
-        pad_img[..., :img.shape[-2], :img.shape[-1]].copy_(img)
+        pad_img[..., : img.shape[-2], : img.shape[-1]].copy_(img)
     return batched_imgs
 
 
@@ -233,10 +228,11 @@ def setup_for_distributed(is_master):
     This function disables printing when not in master process
     """
     import builtins as __builtin__
+
     builtin_print = __builtin__.print
 
     def print(*args, **kwargs):
-        force = kwargs.pop('force', False)
+        force = kwargs.pop("force", False)
         if is_master or force:
             builtin_print(*args, **kwargs)
 
@@ -273,26 +269,26 @@ def save_on_master(*args, **kwargs):
 
 
 def init_distributed_mode(args):
-    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
+    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
         args.rank = int(os.environ["RANK"])
-        args.world_size = int(os.environ['WORLD_SIZE'])
-        args.gpu = int(os.environ['LOCAL_RANK'])
-    elif 'SLURM_PROCID' in os.environ:
-        args.rank = int(os.environ['SLURM_PROCID'])
+        args.world_size = int(os.environ["WORLD_SIZE"])
+        args.gpu = int(os.environ["LOCAL_RANK"])
+    elif "SLURM_PROCID" in os.environ:
+        args.rank = int(os.environ["SLURM_PROCID"])
         args.gpu = args.rank % torch.cuda.device_count()
     elif hasattr(args, "rank"):
         pass
     else:
-        print('Not using distributed mode')
+        print("Not using distributed mode")
         args.distributed = False
         return
 
     args.distributed = True
 
     torch.cuda.set_device(args.gpu)
-    args.dist_backend = 'nccl'
-    print('| distributed init (rank {}): {}'.format(
-        args.rank, args.dist_url), flush=True)
-    torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
-                                         world_size=args.world_size, rank=args.rank)
+    args.dist_backend = "nccl"
+    print("| distributed init (rank {}): {}".format(args.rank, args.dist_url), flush=True)
+    torch.distributed.init_process_group(
+        backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank
+    )
     setup_for_distributed(args.rank == 0)

references/similarity/loss.py

-'''
+"""
     Pytorch adaptation of https://omoindrot.github.io/triplet-loss
     https://github.com/omoindrot/tensorflow-triplet-loss
-'''
+"""
 import torch
 import torch.nn as nn
 
 
 class TripletMarginLoss(nn.Module):
-    def __init__(self, margin=1.0, p=2., mining='batch_all'):
+    def __init__(self, margin=1.0, p=2.0, mining="batch_all"):
         super(TripletMarginLoss, self).__init__()
         self.margin = margin
         self.p = p
         self.mining = mining
 
-        if mining == 'batch_all':
+        if mining == "batch_all":
             self.loss_fn = batch_all_triplet_loss
-        if mining == 'batch_hard':
+        if mining == "batch_hard":
             self.loss_fn = batch_hard_triplet_loss
 
     def forward(self, embeddings, labels):

references/similarity/sampler.py

+import random
+from collections import defaultdict
+
 import torch
 from torch.utils.data.sampler import Sampler
-from collections import defaultdict
-import random
 
 
 def create_groups(groups, k):

references/similarity/test.py

 import unittest
 from collections import defaultdict
 
-from torch.utils.data import DataLoader
-from torchvision.datasets import FakeData
 import torchvision.transforms as transforms
-
 from sampler import PKSampler
+from torch.utils.data import DataLoader
+from torchvision.datasets import FakeData
 
 
 class Tester(unittest.TestCase):
-
     def test_pksampler(self):
         p, k = 16, 4
 
@@ -19,8 +17,7 @@ class Tester(unittest.TestCase):
         self.assertRaises(AssertionError, PKSampler, targets, p, k)
 
         # Ensure p, k constraints on batch
-        dataset = FakeData(size=1000, num_classes=100, image_size=(3, 1, 1),
-                           transform=transforms.ToTensor())
+        dataset = FakeData(size=1000, num_classes=100, image_size=(3, 1, 1), transform=transforms.ToTensor())
         targets = [target.item() for _, target in dataset]
         sampler = PKSampler(targets, p, k)
         loader = DataLoader(dataset, batch_size=p * k, sampler=sampler)
@@ -38,5 +35,5 @@ class Tester(unittest.TestCase):
                 self.assertEqual(bins[b], k)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()

references/similarity/train.py

 import os
 
 import torch
-from torch.optim import Adam
-from torch.utils.data import DataLoader
-
 import torchvision.transforms as transforms
-from torchvision.datasets import FashionMNIST
-
 from loss import TripletMarginLoss
-from sampler import PKSampler
 from model import EmbeddingNet
+from sampler import PKSampler
+from torch.optim import Adam
+from torch.utils.data import DataLoader
+from torchvision.datasets import FashionMNIST
 
 
 def train_epoch(model, optimizer, criterion, data_loader, device, epoch, print_freq):
@@ -33,7 +31,7 @@ def train_epoch(model, optimizer, criterion, data_loader, device, epoch, print_f
             i += 1
             avg_loss = running_loss / print_freq
             avg_trip = 100.0 * running_frac_pos_triplets / print_freq
-            print('[{:d}, {:d}] | loss: {:.4f} | % avg hard triplets: {:.2f}%'.format(epoch, i, avg_loss, avg_trip))
+            print("[{:d}, {:d}] | loss: {:.4f} | % avg hard triplets: {:.2f}%".format(epoch, i, avg_loss, avg_trip))
             running_loss = 0
             running_frac_pos_triplets = 0
 
@@ -79,17 +77,17 @@ def evaluate(model, loader, device):
 
     threshold, accuracy = find_best_threshold(dists, targets, device)
 
-    print('accuracy: {:.3f}%, threshold: {:.2f}'.format(accuracy, threshold))
+    print("accuracy: {:.3f}%, threshold: {:.2f}".format(accuracy, threshold))
 
 
 def save(model, epoch, save_dir, file_name):
-    file_name = 'epoch_' + str(epoch) + '__' + file_name
+    file_name = "epoch_" + str(epoch) + "__" + file_name
     save_path = os.path.join(save_dir, file_name)
     torch.save(model.state_dict(), save_path)
 
 
 def main(args):
-    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     p = args.labels_per_batch
     k = args.samples_per_label
     batch_size = p * k
@@ -103,9 +101,9 @@ def main(args):
     criterion = TripletMarginLoss(margin=args.margin)
     optimizer = Adam(model.parameters(), lr=args.lr)
 
-    transform = transforms.Compose([transforms.Lambda(lambda image: image.convert('RGB')),
-                                    transforms.Resize((224, 224)),
-                                    transforms.ToTensor()])
+    transform = transforms.Compose(
+        [transforms.Lambda(lambda image: image.convert("RGB")), transforms.Resize((224, 224)), transforms.ToTensor()]
+    )
 
     # Using FMNIST to demonstrate embedding learning using triplet loss. This dataset can
     # be replaced with any classification dataset.
@@ -118,48 +116,44 @@ def main(args):
     # targets attribute with the same format.
     targets = train_dataset.targets.tolist()
 
-    train_loader = DataLoader(train_dataset, batch_size=batch_size,
-                              sampler=PKSampler(targets, p, k),
-                              num_workers=args.workers)
-    test_loader = DataLoader(test_dataset, batch_size=args.eval_batch_size,
-                             shuffle=False,
-                             num_workers=args.workers)
+    train_loader = DataLoader(
+        train_dataset, batch_size=batch_size, sampler=PKSampler(targets, p, k), num_workers=args.workers
+    )
+    test_loader = DataLoader(test_dataset, batch_size=args.eval_batch_size, shuffle=False, num_workers=args.workers)
 
     for epoch in range(1, args.epochs + 1):
-        print('Training...')
+        print("Training...")
         train_epoch(model, optimizer, criterion, train_loader, device, epoch, args.print_freq)
 
-        print('Evaluating...')
+        print("Evaluating...")
         evaluate(model, test_loader, device)
 
-        print('Saving...')
-        save(model, epoch, args.save_dir, 'ckpt.pth')
+        print("Saving...")
+        save(model, epoch, args.save_dir, "ckpt.pth")
 
 
 def parse_args():
     import argparse
-    parser = argparse.ArgumentParser(description='PyTorch Embedding Learning')
-
-    parser.add_argument('--dataset-dir', default='/tmp/fmnist/',
-                        help='FashionMNIST dataset directory path')
-    parser.add_argument('-p', '--labels-per-batch', default=8, type=int,
-                        help='Number of unique labels/classes per batch')
-    parser.add_argument('-k', '--samples-per-label', default=8, type=int,
-                        help='Number of samples per label in a batch')
-    parser.add_argument('--eval-batch-size', default=512, type=int)
-    parser.add_argument('--epochs', default=10, type=int, metavar='N',
-                        help='Number of training epochs to run')
-    parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
-                        help='Number of data loading workers')
-    parser.add_argument('--lr', default=0.0001, type=float, help='Learning rate')
-    parser.add_argument('--margin', default=0.2, type=float, help='Triplet loss margin')
-    parser.add_argument('--print-freq', default=20, type=int, help='Print frequency')
-    parser.add_argument('--save-dir', default='.', help='Model save directory')
-    parser.add_argument('--resume', default='', help='Resume from checkpoint')
+
+    parser = argparse.ArgumentParser(description="PyTorch Embedding Learning")
+
+    parser.add_argument("--dataset-dir", default="/tmp/fmnist/", help="FashionMNIST dataset directory path")
+    parser.add_argument(
+        "-p", "--labels-per-batch", default=8, type=int, help="Number of unique labels/classes per batch"
+    )
+    parser.add_argument("-k", "--samples-per-label", default=8, type=int, help="Number of samples per label in a batch")
+    parser.add_argument("--eval-batch-size", default=512, type=int)
+    parser.add_argument("--epochs", default=10, type=int, metavar="N", help="Number of training epochs to run")
+    parser.add_argument("-j", "--workers", default=4, type=int, metavar="N", help="Number of data loading workers")
+    parser.add_argument("--lr", default=0.0001, type=float, help="Learning rate")
+    parser.add_argument("--margin", default=0.2, type=float, help="Triplet loss margin")
+    parser.add_argument("--print-freq", default=20, type=int, help="Print frequency")
+    parser.add_argument("--save-dir", default=".", help="Model save directory")
+    parser.add_argument("--resume", default="", help="Resume from checkpoint")
 
     return parser.parse_args()
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     args = parse_args()
     main(args)

references/video_classification/presets.py

 import torch
-
 from torchvision.transforms import transforms
 from transforms import ConvertBHWCtoBCHW, ConvertBCHWtoCBHW
 
 
 class VideoClassificationPresetTrain:
-    def __init__(self, resize_size, crop_size, mean=(0.43216, 0.394666, 0.37645), std=(0.22803, 0.22145, 0.216989),
-                 hflip_prob=0.5):
+    def __init__(
+        self,
+        resize_size,
+        crop_size,
+        mean=(0.43216, 0.394666, 0.37645),
+        std=(0.22803, 0.22145, 0.216989),
+        hflip_prob=0.5,
+    ):
         trans = [
             ConvertBHWCtoBCHW(),
             transforms.ConvertImageDtype(torch.float32),
@@ -14,11 +19,7 @@ class VideoClassificationPresetTrain:
         ]
         if hflip_prob > 0:
             trans.append(transforms.RandomHorizontalFlip(hflip_prob))
-        trans.extend([
-            transforms.Normalize(mean=mean, std=std),
-            transforms.RandomCrop(crop_size),
-            ConvertBCHWtoCBHW()
-        ])
+        trans.extend([transforms.Normalize(mean=mean, std=std), transforms.RandomCrop(crop_size), ConvertBCHWtoCBHW()])
         self.transforms = transforms.Compose(trans)
 
     def __call__(self, x):
@@ -27,14 +28,16 @@ class VideoClassificationPresetTrain:
 
 class VideoClassificationPresetEval:
     def __init__(self, resize_size, crop_size, mean=(0.43216, 0.394666, 0.37645), std=(0.22803, 0.22145, 0.216989)):
-        self.transforms = transforms.Compose([
-            ConvertBHWCtoBCHW(),
-            transforms.ConvertImageDtype(torch.float32),
-            transforms.Resize(resize_size),
-            transforms.Normalize(mean=mean, std=std),
-            transforms.CenterCrop(crop_size),
-            ConvertBCHWtoCBHW()
-        ])
+        self.transforms = transforms.Compose(
+            [
+                ConvertBHWCtoBCHW(),
+                transforms.ConvertImageDtype(torch.float32),
+                transforms.Resize(resize_size),
+                transforms.Normalize(mean=mean, std=std),
+                transforms.CenterCrop(crop_size),
+                ConvertBCHWtoCBHW(),
+            ]
+        )
 
     def __call__(self, x):
         return self.transforms(x)

references/video_classification/train.py

 import datetime
 import os
 import time
+
+import presets
 import torch
 import torch.utils.data
-from torch.utils.data.dataloader import default_collate
-from torch import nn
 import torchvision
 import torchvision.datasets.video_utils
-from torchvision.datasets.samplers import DistributedSampler, UniformClipSampler, RandomClipSampler
-
-import presets
 import utils
+from torch import nn
+from torch.utils.data.dataloader import default_collate
+from torchvision.datasets.samplers import DistributedSampler, UniformClipSampler, RandomClipSampler
 
 try:
     from apex import amp
@@ -21,10 +21,10 @@ except ImportError:
 def train_one_epoch(model, criterion, optimizer, lr_scheduler, data_loader, device, epoch, print_freq, apex=False):
     model.train()
     metric_logger = utils.MetricLogger(delimiter="  ")
-    metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value}'))
-    metric_logger.add_meter('clips/s', utils.SmoothedValue(window_size=10, fmt='{value:.3f}'))
+    metric_logger.add_meter("lr", utils.SmoothedValue(window_size=1, fmt="{value}"))
+    metric_logger.add_meter("clips/s", utils.SmoothedValue(window_size=10, fmt="{value:.3f}"))
 
-    header = 'Epoch: [{}]'.format(epoch)
+    header = "Epoch: [{}]".format(epoch)
     for video, target in metric_logger.log_every(data_loader, print_freq, header):
         start_time = time.time()
         video, target = video.to(device), target.to(device)
@@ -42,16 +42,16 @@ def train_one_epoch(model, criterion, optimizer, lr_scheduler, data_loader, devi
         acc1, acc5 = utils.accuracy(output, target, topk=(1, 5))
         batch_size = video.shape[0]
         metric_logger.update(loss=loss.item(), lr=optimizer.param_groups[0]["lr"])
-        metric_logger.meters['acc1'].update(acc1.item(), n=batch_size)
-        metric_logger.meters['acc5'].update(acc5.item(), n=batch_size)
-        metric_logger.meters['clips/s'].update(batch_size / (time.time() - start_time))
+        metric_logger.meters["acc1"].update(acc1.item(), n=batch_size)
+        metric_logger.meters["acc5"].update(acc5.item(), n=batch_size)
+        metric_logger.meters["clips/s"].update(batch_size / (time.time() - start_time))
         lr_scheduler.step()
 
 
 def evaluate(model, criterion, data_loader, device):
     model.eval()
     metric_logger = utils.MetricLogger(delimiter="  ")
-    header = 'Test:'
+    header = "Test:"
     with torch.no_grad():
         for video, target in metric_logger.log_every(data_loader, 100, header):
             video = video.to(device, non_blocking=True)
@@ -64,18 +64,22 @@ def evaluate(model, criterion, data_loader, device):
             # could have been padded in distributed setup
             batch_size = video.shape[0]
             metric_logger.update(loss=loss.item())
-            metric_logger.meters['acc1'].update(acc1.item(), n=batch_size)
-            metric_logger.meters['acc5'].update(acc5.item(), n=batch_size)
+            metric_logger.meters["acc1"].update(acc1.item(), n=batch_size)
+            metric_logger.meters["acc5"].update(acc5.item(), n=batch_size)
     # gather the stats from all processes
     metric_logger.synchronize_between_processes()
 
-    print(' * Clip Acc@1 {top1.global_avg:.3f} Clip Acc@5 {top5.global_avg:.3f}'
-          .format(top1=metric_logger.acc1, top5=metric_logger.acc5))
+    print(
+        " * Clip Acc@1 {top1.global_avg:.3f} Clip Acc@5 {top5.global_avg:.3f}".format(
+            top1=metric_logger.acc1, top5=metric_logger.acc5
+        )
+    )
     return metric_logger.acc1.global_avg
 
 
 def _get_cache_path(filepath):
     import hashlib
+
     h = hashlib.sha1(filepath.encode()).hexdigest()
     cache_path = os.path.join("~", ".torch", "vision", "datasets", "kinetics", h[:10] + ".pt")
     cache_path = os.path.expanduser(cache_path)
@@ -90,8 +94,10 @@ def collate_fn(batch):
 
 def main(args):
     if args.apex and amp is None:
-        raise RuntimeError("Failed to import apex. Please install apex from https://www.github.com/nvidia/apex "
-                           "to enable mixed-precision training.")
+        raise RuntimeError(
+            "Failed to import apex. Please install apex from https://www.github.com/nvidia/apex "
+            "to enable mixed-precision training."
+        )
 
     if args.output_dir:
         utils.mkdir(args.output_dir)
@@ -121,15 +127,17 @@ def main(args):
         dataset.transform = transform_train
     else:
         if args.distributed:
-            print("It is recommended to pre-compute the dataset cache "
-                  "on a single-gpu first, as it will be faster")
+            print("It is recommended to pre-compute the dataset cache " "on a single-gpu first, as it will be faster")
         dataset = torchvision.datasets.Kinetics400(
             traindir,
             frames_per_clip=args.clip_len,
             step_between_clips=1,
             transform=transform_train,
             frame_rate=15,
-            extensions=('avi', 'mp4', )
+            extensions=(
+                "avi",
+                "mp4",
+            ),
         )
         if args.cache_dataset:
             print("Saving dataset_train to {}".format(cache_path))
@@ -149,15 +157,17 @@ def main(args):
         dataset_test.transform = transform_test
     else:
         if args.distributed:
-            print("It is recommended to pre-compute the dataset cache "
-                  "on a single-gpu first, as it will be faster")
+            print("It is recommended to pre-compute the dataset cache " "on a single-gpu first, as it will be faster")
         dataset_test = torchvision.datasets.Kinetics400(
             valdir,
             frames_per_clip=args.clip_len,
             step_between_clips=1,
             transform=transform_test,
             frame_rate=15,
-            extensions=('avi', 'mp4',)
+            extensions=(
+                "avi",
+                "mp4",
+            ),
         )
         if args.cache_dataset:
             print("Saving dataset_test to {}".format(cache_path))
@@ -172,14 +182,22 @@ def main(args):
         test_sampler = DistributedSampler(test_sampler)
 
     data_loader = torch.utils.data.DataLoader(
-        dataset, batch_size=args.batch_size,
-        sampler=train_sampler, num_workers=args.workers,
-        pin_memory=True, collate_fn=collate_fn)
+        dataset,
+        batch_size=args.batch_size,
+        sampler=train_sampler,
+        num_workers=args.workers,
+        pin_memory=True,
+        collate_fn=collate_fn,
+    )
 
     data_loader_test = torch.utils.data.DataLoader(
-        dataset_test, batch_size=args.batch_size,
-        sampler=test_sampler, num_workers=args.workers,
-        pin_memory=True, collate_fn=collate_fn)
+        dataset_test,
+        batch_size=args.batch_size,
+        sampler=test_sampler,
+        num_workers=args.workers,
+        pin_memory=True,
+        collate_fn=collate_fn,
+    )
 
     print("Creating model")
     model = torchvision.models.video.__dict__[args.model](pretrained=args.pretrained)
@@ -190,13 +208,10 @@ def main(args):
     criterion = nn.CrossEntropyLoss()
 
     lr = args.lr * args.world_size
-    optimizer = torch.optim.SGD(
-        model.parameters(), lr=lr, momentum=args.momentum, weight_decay=args.weight_decay)
+    optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=args.momentum, weight_decay=args.weight_decay)
 
     if args.apex:
-        model, optimizer = amp.initialize(model, optimizer,
-                                          opt_level=args.apex_opt_level
-                                          )
+        model, optimizer = amp.initialize(model, optimizer, opt_level=args.apex_opt_level)
 
     # convert scheduler to be per iteration, not per epoch, for warmup that lasts
     # between different epochs
@@ -207,20 +222,22 @@ def main(args):
     if args.lr_warmup_epochs > 0:
         warmup_iters = iters_per_epoch * args.lr_warmup_epochs
         args.lr_warmup_method = args.lr_warmup_method.lower()
-        if args.lr_warmup_method == 'linear':
-            warmup_lr_scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=args.lr_warmup_decay,
-                                                                    total_iters=warmup_iters)
-        elif args.lr_warmup_method == 'constant':
-            warmup_lr_scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer, factor=args.lr_warmup_decay,
-                                                                      total_iters=warmup_iters)
+        if args.lr_warmup_method == "linear":
+            warmup_lr_scheduler = torch.optim.lr_scheduler.LinearLR(
+                optimizer, start_factor=args.lr_warmup_decay, total_iters=warmup_iters
+            )
+        elif args.lr_warmup_method == "constant":
+            warmup_lr_scheduler = torch.optim.lr_scheduler.ConstantLR(
+                optimizer, factor=args.lr_warmup_decay, total_iters=warmup_iters
+            )
         else:
-            raise RuntimeError("Invalid warmup lr method '{}'. Only linear and constant "
-                               "are supported.".format(args.lr_warmup_method))
+            raise RuntimeError(
+                "Invalid warmup lr method '{}'. Only linear and constant "
+                "are supported.".format(args.lr_warmup_method)
+            )
 
         lr_scheduler = torch.optim.lr_scheduler.SequentialLR(
-            optimizer,
-            schedulers=[warmup_lr_scheduler, main_lr_scheduler],
-            milestones=[warmup_iters]
+            optimizer, schedulers=[warmup_lr_scheduler, main_lr_scheduler], milestones=[warmup_iters]
         )
     else:
         lr_scheduler = main_lr_scheduler
@@ -231,11 +248,11 @@ def main(args):
         model_without_ddp = model.module
 
     if args.resume:
-        checkpoint = torch.load(args.resume, map_location='cpu')
-        model_without_ddp.load_state_dict(checkpoint['model'])
-        optimizer.load_state_dict(checkpoint['optimizer'])
-        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
-        args.start_epoch = checkpoint['epoch'] + 1
+        checkpoint = torch.load(args.resume, map_location="cpu")
+        model_without_ddp.load_state_dict(checkpoint["model"])
+        optimizer.load_state_dict(checkpoint["optimizer"])
+        lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
+        args.start_epoch = checkpoint["epoch"] + 1
 
     if args.test_only:
         evaluate(model, criterion, data_loader_test, device=device)
@@ -246,62 +263,65 @@ def main(args):
     for epoch in range(args.start_epoch, args.epochs):
         if args.distributed:
             train_sampler.set_epoch(epoch)
-        train_one_epoch(model, criterion, optimizer, lr_scheduler, data_loader,
-                        device, epoch, args.print_freq, args.apex)
+        train_one_epoch(
+            model, criterion, optimizer, lr_scheduler, data_loader, device, epoch, args.print_freq, args.apex
+        )
         evaluate(model, criterion, data_loader_test, device=device)
         if args.output_dir:
             checkpoint = {
-                'model': model_without_ddp.state_dict(),
-                'optimizer': optimizer.state_dict(),
-                'lr_scheduler': lr_scheduler.state_dict(),
-                'epoch': epoch,
-                'args': args}
-            utils.save_on_master(
-                checkpoint,
-                os.path.join(args.output_dir, 'model_{}.pth'.format(epoch)))
-            utils.save_on_master(
-                checkpoint,
-                os.path.join(args.output_dir, 'checkpoint.pth'))
+                "model": model_without_ddp.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "lr_scheduler": lr_scheduler.state_dict(),
+                "epoch": epoch,
+                "args": args,
+            }
+            utils.save_on_master(checkpoint, os.path.join(args.output_dir, "model_{}.pth".format(epoch)))
+            utils.save_on_master(checkpoint, os.path.join(args.output_dir, "checkpoint.pth"))
 
     total_time = time.time() - start_time
     total_time_str = str(datetime.timedelta(seconds=int(total_time)))
-    print('Training time {}'.format(total_time_str))
+    print("Training time {}".format(total_time_str))
 
 
 def parse_args():
     import argparse
-    parser = argparse.ArgumentParser(description='PyTorch Video Classification Training')
-
-    parser.add_argument('--data-path', default='/datasets01_101/kinetics/070618/', help='dataset')
-    parser.add_argument('--train-dir', default='train_avi-480p', help='name of train dir')
-    parser.add_argument('--val-dir', default='val_avi-480p', help='name of val dir')
-    parser.add_argument('--model', default='r2plus1d_18', help='model')
-    parser.add_argument('--device', default='cuda', help='device')
-    parser.add_argument('--clip-len', default=16, type=int, metavar='N',
-                        help='number of frames per clip')
-    parser.add_argument('--clips-per-video', default=5, type=int, metavar='N',
-                        help='maximum number of clips per video to consider')
-    parser.add_argument('-b', '--batch-size', default=24, type=int)
-    parser.add_argument('--epochs', default=45, type=int, metavar='N',
-                        help='number of total epochs to run')
-    parser.add_argument('-j', '--workers', default=10, type=int, metavar='N',
-                        help='number of data loading workers (default: 10)')
-    parser.add_argument('--lr', default=0.01, type=float, help='initial learning rate')
-    parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
-                        help='momentum')
-    parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
-                        metavar='W', help='weight decay (default: 1e-4)',
-                        dest='weight_decay')
-    parser.add_argument('--lr-milestones', nargs='+', default=[20, 30, 40], type=int, help='decrease lr on milestones')
-    parser.add_argument('--lr-gamma', default=0.1, type=float, help='decrease lr by a factor of lr-gamma')
-    parser.add_argument('--lr-warmup-epochs', default=10, type=int, help='the number of epochs to warmup (default: 10)')
-    parser.add_argument('--lr-warmup-method', default="linear", type=str, help='the warmup method (default: linear)')
-    parser.add_argument('--lr-warmup-decay', default=0.001, type=float, help='the decay for lr')
-    parser.add_argument('--print-freq', default=10, type=int, help='print frequency')
-    parser.add_argument('--output-dir', default='.', help='path where to save')
-    parser.add_argument('--resume', default='', help='resume from checkpoint')
-    parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
-                        help='start epoch')
+
+    parser = argparse.ArgumentParser(description="PyTorch Video Classification Training")
+
+    parser.add_argument("--data-path", default="/datasets01_101/kinetics/070618/", help="dataset")
+    parser.add_argument("--train-dir", default="train_avi-480p", help="name of train dir")
+    parser.add_argument("--val-dir", default="val_avi-480p", help="name of val dir")
+    parser.add_argument("--model", default="r2plus1d_18", help="model")
+    parser.add_argument("--device", default="cuda", help="device")
+    parser.add_argument("--clip-len", default=16, type=int, metavar="N", help="number of frames per clip")
+    parser.add_argument(
+        "--clips-per-video", default=5, type=int, metavar="N", help="maximum number of clips per video to consider"
+    )
+    parser.add_argument("-b", "--batch-size", default=24, type=int)
+    parser.add_argument("--epochs", default=45, type=int, metavar="N", help="number of total epochs to run")
+    parser.add_argument(
+        "-j", "--workers", default=10, type=int, metavar="N", help="number of data loading workers (default: 10)"
+    )
+    parser.add_argument("--lr", default=0.01, type=float, help="initial learning rate")
+    parser.add_argument("--momentum", default=0.9, type=float, metavar="M", help="momentum")
+    parser.add_argument(
+        "--wd",
+        "--weight-decay",
+        default=1e-4,
+        type=float,
+        metavar="W",
+        help="weight decay (default: 1e-4)",
+        dest="weight_decay",
+    )
+    parser.add_argument("--lr-milestones", nargs="+", default=[20, 30, 40], type=int, help="decrease lr on milestones")
+    parser.add_argument("--lr-gamma", default=0.1, type=float, help="decrease lr by a factor of lr-gamma")
+    parser.add_argument("--lr-warmup-epochs", default=10, type=int, help="the number of epochs to warmup (default: 10)")
+    parser.add_argument("--lr-warmup-method", default="linear", type=str, help="the warmup method (default: linear)")
+    parser.add_argument("--lr-warmup-decay", default=0.001, type=float, help="the decay for lr")
+    parser.add_argument("--print-freq", default=10, type=int, help="print frequency")
+    parser.add_argument("--output-dir", default=".", help="path where to save")
+    parser.add_argument("--resume", default="", help="resume from checkpoint")
+    parser.add_argument("--start-epoch", default=0, type=int, metavar="N", help="start epoch")
     parser.add_argument(
         "--cache-dataset",
         dest="cache_dataset",
@@ -328,18 +348,19 @@ def parse_args():
     )
 
     # Mixed precision training parameters
-    parser.add_argument('--apex', action='store_true',
-                        help='Use apex for mixed precision training')
-    parser.add_argument('--apex-opt-level', default='O1', type=str,
-                        help='For apex mixed precision training'
-                             'O0 for FP32 training, O1 for mixed precision training.'
-                             'For further detail, see https://github.com/NVIDIA/apex/tree/master/examples/imagenet'
-                        )
+    parser.add_argument("--apex", action="store_true", help="Use apex for mixed precision training")
+    parser.add_argument(
+        "--apex-opt-level",
+        default="O1",
+        type=str,
+        help="For apex mixed precision training"
+        "O0 for FP32 training, O1 for mixed precision training."
+        "For further detail, see https://github.com/NVIDIA/apex/tree/master/examples/imagenet",
+    )
 
     # distributed training parameters
-    parser.add_argument('--world-size', default=1, type=int,
-                        help='number of distributed processes')
-    parser.add_argument('--dist-url', default='env://', help='url used to set up distributed training')
+    parser.add_argument("--world-size", default=1, type=int, help="number of distributed processes")
+    parser.add_argument("--dist-url", default="env://", help="url used to set up distributed training")
 
     args = parser.parse_args()
 

references/video_classification/transforms.py

@@ -3,16 +3,14 @@ import torch.nn as nn
 
 
 class ConvertBHWCtoBCHW(nn.Module):
-    """Convert tensor from (B, H, W, C) to (B, C, H, W)
-    """
+    """Convert tensor from (B, H, W, C) to (B, C, H, W)"""
 
     def forward(self, vid: torch.Tensor) -> torch.Tensor:
         return vid.permute(0, 3, 1, 2)
 
 
 class ConvertBCHWtoCBHW(nn.Module):
-    """Convert tensor from (B, C, H, W) to (C, B, H, W)
-    """
+    """Convert tensor from (B, C, H, W) to (C, B, H, W)"""
 
     def forward(self, vid: torch.Tensor) -> torch.Tensor:
         return vid.permute(1, 0, 2, 3)

references/video_classification/utils.py

-from collections import defaultdict, deque
 import datetime
+import errno
+import os
 import time
+from collections import defaultdict, deque
+
 import torch
 import torch.distributed as dist
 
-import errno
-import os
-
 
 class SmoothedValue(object):
     """Track a series of values and provide access to smoothed values over a
@@ -32,7 +32,7 @@ class SmoothedValue(object):
         """
         if not is_dist_avail_and_initialized():
             return
-        t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda")
         dist.barrier()
         dist.all_reduce(t)
         t = t.tolist()
@@ -63,11 +63,8 @@ class SmoothedValue(object):
 
     def __str__(self):
         return self.fmt.format(
-            median=self.median,
-            avg=self.avg,
-            global_avg=self.global_avg,
-            max=self.max,
-            value=self.value)
+            median=self.median, avg=self.avg, global_avg=self.global_avg, max=self.max, value=self.value
+        )
 
 
 class MetricLogger(object):
@@ -87,15 +84,12 @@ class MetricLogger(object):
             return self.meters[attr]
         if attr in self.__dict__:
             return self.__dict__[attr]
-        raise AttributeError("'{}' object has no attribute '{}'".format(
-            type(self).__name__, attr))
+        raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, attr))
 
     def __str__(self):
         loss_str = []
         for name, meter in self.meters.items():
-            loss_str.append(
-                "{}: {}".format(name, str(meter))
-            )
+            loss_str.append("{}: {}".format(name, str(meter)))
         return self.delimiter.join(loss_str)
 
     def synchronize_between_processes(self):
@@ -108,31 +102,28 @@ class MetricLogger(object):
     def log_every(self, iterable, print_freq, header=None):
         i = 0
         if not header:
-            header = ''
+            header = ""
         start_time = time.time()
         end = time.time()
-        iter_time = SmoothedValue(fmt='{avg:.4f}')
-        data_time = SmoothedValue(fmt='{avg:.4f}')
-        space_fmt = ':' + str(len(str(len(iterable)))) + 'd'
+        iter_time = SmoothedValue(fmt="{avg:.4f}")
+        data_time = SmoothedValue(fmt="{avg:.4f}")
+        space_fmt = ":" + str(len(str(len(iterable)))) + "d"
         if torch.cuda.is_available():
-            log_msg = self.delimiter.join([
-                header,
-                '[{0' + space_fmt + '}/{1}]',
-                'eta: {eta}',
-                '{meters}',
-                'time: {time}',
-                'data: {data}',
-                'max mem: {memory:.0f}'
-            ])
+            log_msg = self.delimiter.join(
+                [
+                    header,
+                    "[{0" + space_fmt + "}/{1}]",
+                    "eta: {eta}",
+                    "{meters}",
+                    "time: {time}",
+                    "data: {data}",
+                    "max mem: {memory:.0f}",
+                ]
+            )
         else:
-            log_msg = self.delimiter.join([
-                header,
-                '[{0' + space_fmt + '}/{1}]',
-                'eta: {eta}',
-                '{meters}',
-                'time: {time}',
-                'data: {data}'
-            ])
+            log_msg = self.delimiter.join(
+                [header, "[{0" + space_fmt + "}/{1}]", "eta: {eta}", "{meters}", "time: {time}", "data: {data}"]
+            )
         MB = 1024.0 * 1024.0
         for obj in iterable:
             data_time.update(time.time() - end)
@@ -142,21 +133,28 @@ class MetricLogger(object):
                 eta_seconds = iter_time.global_avg * (len(iterable) - i)
                 eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
                 if torch.cuda.is_available():
-                    print(log_msg.format(
-                        i, len(iterable), eta=eta_string,
-                        meters=str(self),
-                        time=str(iter_time), data=str(data_time),
-                        memory=torch.cuda.max_memory_allocated() / MB))
+                    print(
+                        log_msg.format(
+                            i,
+                            len(iterable),
+                            eta=eta_string,
+                            meters=str(self),
+                            time=str(iter_time),
+                            data=str(data_time),
+                            memory=torch.cuda.max_memory_allocated() / MB,
+                        )
+                    )
                 else:
-                    print(log_msg.format(
-                        i, len(iterable), eta=eta_string,
-                        meters=str(self),
-                        time=str(iter_time), data=str(data_time)))
+                    print(
+                        log_msg.format(
+                            i, len(iterable), eta=eta_string, meters=str(self), time=str(iter_time), data=str(data_time)
+                        )
+                    )
             i += 1
             end = time.time()
         total_time = time.time() - start_time
         total_time_str = str(datetime.timedelta(seconds=int(total_time)))
-        print('{} Total time: {}'.format(header, total_time_str))
+        print("{} Total time: {}".format(header, total_time_str))
 
 
 def accuracy(output, target, topk=(1,)):
@@ -189,10 +187,11 @@ def setup_for_distributed(is_master):
     This function disables printing when not in master process
     """
     import builtins as __builtin__
+
     builtin_print = __builtin__.print
 
     def print(*args, **kwargs):
-        force = kwargs.pop('force', False)
+        force = kwargs.pop("force", False)
         if is_master or force:
             builtin_print(*args, **kwargs)
 
@@ -229,26 +228,26 @@ def save_on_master(*args, **kwargs):
 
 
 def init_distributed_mode(args):
-    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
+    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
         args.rank = int(os.environ["RANK"])
-        args.world_size = int(os.environ['WORLD_SIZE'])
-        args.gpu = int(os.environ['LOCAL_RANK'])
-    elif 'SLURM_PROCID' in os.environ:
-        args.rank = int(os.environ['SLURM_PROCID'])
+        args.world_size = int(os.environ["WORLD_SIZE"])
+        args.gpu = int(os.environ["LOCAL_RANK"])
+    elif "SLURM_PROCID" in os.environ:
+        args.rank = int(os.environ["SLURM_PROCID"])
         args.gpu = args.rank % torch.cuda.device_count()
     elif hasattr(args, "rank"):
         pass
     else:
-        print('Not using distributed mode')
+        print("Not using distributed mode")
         args.distributed = False
         return
 
     args.distributed = True
 
     torch.cuda.set_device(args.gpu)
-    args.dist_backend = 'nccl'
-    print('| distributed init (rank {}): {}'.format(
-        args.rank, args.dist_url), flush=True)
-    torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
-                                         world_size=args.world_size, rank=args.rank)
+    args.dist_backend = "nccl"
+    print("| distributed init (rank {}): {}".format(args.rank, args.dist_url), flush=True)
+    torch.distributed.init_process_group(
+        backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank
+    )
     setup_for_distributed(args.rank == 0)

setup.cfg

@@ -9,7 +9,13 @@ max-line-length = 120
 
 [flake8]
 max-line-length = 120
-ignore = F401,E402,F403,W503,W504,F821
+ignore = E203, E402, W503, W504, F821
+per-file-ignores =
+    __init__.py: F401, F403, F405
+    ./hubconf.py: F401
+    torchvision/models/mobilenet.py: F401, F403
+    torchvision/models/quantization/mobilenet.py: F401, F403
+    test/smoke_test.py: F401
 exclude = venv
 
 [pydocstyle]

setup.py

-import os
-import io
-import re
-import sys
-from setuptools import setup, find_packages
-from pkg_resources import parse_version, get_distribution, DistributionNotFound
-import subprocess
 import distutils.command.clean
 import distutils.spawn
-from distutils.version import StrictVersion
 import glob
+import io
+import os
+import re
 import shutil
+import subprocess
+import sys
+from distutils.version import StrictVersion
 
 import torch
+from pkg_resources import parse_version, get_distribution, DistributionNotFound
+from setuptools import setup, find_packages
 from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension, CUDA_HOME
 
 
 def read(*names, **kwargs):
-    with io.open(
-        os.path.join(os.path.dirname(__file__), *names),
-        encoding=kwargs.get("encoding", "utf8")
-    ) as fp:
+    with io.open(os.path.join(os.path.dirname(__file__), *names), encoding=kwargs.get("encoding", "utf8")) as fp:
         return fp.read()
 
 
@@ -32,26 +29,26 @@ def get_dist(pkgname):
 
 cwd = os.path.dirname(os.path.abspath(__file__))
 
-version_txt = os.path.join(cwd, 'version.txt')
-with open(version_txt, 'r') as f:
+version_txt = os.path.join(cwd, "version.txt")
+with open(version_txt, "r") as f:
     version = f.readline().strip()
-sha = 'Unknown'
-package_name = 'torchvision'
+sha = "Unknown"
+package_name = "torchvision"
 
 try:
-    sha = subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=cwd).decode('ascii').strip()
+    sha = subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=cwd).decode("ascii").strip()
 except Exception:
     pass
 
-if os.getenv('BUILD_VERSION'):
-    version = os.getenv('BUILD_VERSION')
-elif sha != 'Unknown':
-    version += '+' + sha[:7]
+if os.getenv("BUILD_VERSION"):
+    version = os.getenv("BUILD_VERSION")
+elif sha != "Unknown":
+    version += "+" + sha[:7]
 
 
 def write_version_file():
-    version_path = os.path.join(cwd, 'torchvision', 'version.py')
-    with open(version_path, 'w') as f:
+    version_path = os.path.join(cwd, "torchvision", "version.py")
+    with open(version_path, "w") as f:
         f.write("__version__ = '{}'\n".format(version))
         f.write("git_version = {}\n".format(repr(sha)))
         f.write("from torchvision.extension import _check_cuda_version\n")
@@ -59,34 +56,34 @@ def write_version_file():
         f.write("    cuda = _check_cuda_version()\n")
 
 
-pytorch_dep = 'torch'
-if os.getenv('PYTORCH_VERSION'):
-    pytorch_dep += "==" + os.getenv('PYTORCH_VERSION')
+pytorch_dep = "torch"
+if os.getenv("PYTORCH_VERSION"):
+    pytorch_dep += "==" + os.getenv("PYTORCH_VERSION")
 
 requirements = [
-    'numpy',
+    "numpy",
     pytorch_dep,
 ]
 
 # Excluding 8.3.0 because of https://github.com/pytorch/vision/issues/4146
-pillow_ver = ' >= 5.3.0, !=8.3.0'
-pillow_req = 'pillow-simd' if get_dist('pillow-simd') is not None else 'pillow'
+pillow_ver = " >= 5.3.0, !=8.3.0"
+pillow_req = "pillow-simd" if get_dist("pillow-simd") is not None else "pillow"
 requirements.append(pillow_req + pillow_ver)
 
 
 def find_library(name, vision_include):
     this_dir = os.path.dirname(os.path.abspath(__file__))
-    build_prefix = os.environ.get('BUILD_PREFIX', None)
+    build_prefix = os.environ.get("BUILD_PREFIX", None)
     is_conda_build = build_prefix is not None
 
     library_found = False
     conda_installed = False
     lib_folder = None
     include_folder = None
-    library_header = '{0}.h'.format(name)
+    library_header = "{0}.h".format(name)
 
     # Lookup in TORCHVISION_INCLUDE or in the package file
-    package_path = [os.path.join(this_dir, 'torchvision')]
+    package_path = [os.path.join(this_dir, "torchvision")]
     for folder in vision_include + package_path:
         candidate_path = os.path.join(folder, library_header)
         library_found = os.path.exists(candidate_path)
@@ -94,67 +91,66 @@ def find_library(name, vision_include):
             break
 
     if not library_found:
-        print('Running build on conda-build: {0}'.format(is_conda_build))
+        print("Running build on conda-build: {0}".format(is_conda_build))
         if is_conda_build:
             # Add conda headers/libraries
-            if os.name == 'nt':
-                build_prefix = os.path.join(build_prefix, 'Library')
-            include_folder = os.path.join(build_prefix, 'include')
-            lib_folder = os.path.join(build_prefix, 'lib')
-            library_header_path = os.path.join(
-                include_folder, library_header)
+            if os.name == "nt":
+                build_prefix = os.path.join(build_prefix, "Library")
+            include_folder = os.path.join(build_prefix, "include")
+            lib_folder = os.path.join(build_prefix, "lib")
+            library_header_path = os.path.join(include_folder, library_header)
             library_found = os.path.isfile(library_header_path)
             conda_installed = library_found
         else:
             # Check if using Anaconda to produce wheels
-            conda = distutils.spawn.find_executable('conda')
+            conda = distutils.spawn.find_executable("conda")
             is_conda = conda is not None
-            print('Running build on conda: {0}'.format(is_conda))
+            print("Running build on conda: {0}".format(is_conda))
             if is_conda:
                 python_executable = sys.executable
                 py_folder = os.path.dirname(python_executable)
-                if os.name == 'nt':
-                    env_path = os.path.join(py_folder, 'Library')
+                if os.name == "nt":
+                    env_path = os.path.join(py_folder, "Library")
                 else:
                     env_path = os.path.dirname(py_folder)
-                lib_folder = os.path.join(env_path, 'lib')
-                include_folder = os.path.join(env_path, 'include')
-                library_header_path = os.path.join(
-                    include_folder, library_header)
+                lib_folder = os.path.join(env_path, "lib")
+                include_folder = os.path.join(env_path, "include")
+                library_header_path = os.path.join(include_folder, library_header)
                 library_found = os.path.isfile(library_header_path)
                 conda_installed = library_found
 
         if not library_found:
-            if sys.platform == 'linux':
-                library_found = os.path.exists('/usr/include/{0}'.format(
-                    library_header))
-                library_found = library_found or os.path.exists(
-                    '/usr/local/include/{0}'.format(library_header))
+            if sys.platform == "linux":
+                library_found = os.path.exists("/usr/include/{0}".format(library_header))
+                library_found = library_found or os.path.exists("/usr/local/include/{0}".format(library_header))
 
     return library_found, conda_installed, include_folder, lib_folder
 
 
 def get_extensions():
     this_dir = os.path.dirname(os.path.abspath(__file__))
-    extensions_dir = os.path.join(this_dir, 'torchvision', 'csrc')
+    extensions_dir = os.path.join(this_dir, "torchvision", "csrc")
 
-    main_file = glob.glob(os.path.join(extensions_dir, '*.cpp')) + glob.glob(os.path.join(extensions_dir, 'ops',
-                                                                                          '*.cpp'))
+    main_file = glob.glob(os.path.join(extensions_dir, "*.cpp")) + glob.glob(
+        os.path.join(extensions_dir, "ops", "*.cpp")
+    )
     source_cpu = (
-        glob.glob(os.path.join(extensions_dir, 'ops', 'autograd', '*.cpp')) +
-        glob.glob(os.path.join(extensions_dir, 'ops', 'cpu', '*.cpp')) +
-        glob.glob(os.path.join(extensions_dir, 'ops', 'quantized', 'cpu', '*.cpp'))
+        glob.glob(os.path.join(extensions_dir, "ops", "autograd", "*.cpp"))
+        + glob.glob(os.path.join(extensions_dir, "ops", "cpu", "*.cpp"))
+        + glob.glob(os.path.join(extensions_dir, "ops", "quantized", "cpu", "*.cpp"))
     )
 
     is_rocm_pytorch = False
-    TORCH_MAJOR = int(torch.__version__.split('.')[0])
-    TORCH_MINOR = int(torch.__version__.split('.')[1])
+    TORCH_MAJOR = int(torch.__version__.split(".")[0])
+    TORCH_MINOR = int(torch.__version__.split(".")[1])
     if TORCH_MAJOR > 1 or (TORCH_MAJOR == 1 and TORCH_MINOR >= 5):
         from torch.utils.cpp_extension import ROCM_HOME
+
         is_rocm_pytorch = True if ((torch.version.hip is not None) and (ROCM_HOME is not None)) else False
 
     if is_rocm_pytorch:
         from torch.utils.hipify import hipify_python
+
         hipify_python.hipify(
             project_directory=this_dir,
             output_directory=this_dir,
@@ -162,25 +158,25 @@ def get_extensions():
             show_detailed=True,
             is_pytorch_extension=True,
         )
-        source_cuda = glob.glob(os.path.join(extensions_dir, 'ops', 'hip', '*.hip'))
+        source_cuda = glob.glob(os.path.join(extensions_dir, "ops", "hip", "*.hip"))
         # Copy over additional files
         for file in glob.glob(r"torchvision/csrc/ops/cuda/*.h"):
             shutil.copy(file, "torchvision/csrc/ops/hip")
 
     else:
-        source_cuda = glob.glob(os.path.join(extensions_dir, 'ops', 'cuda', '*.cu'))
+        source_cuda = glob.glob(os.path.join(extensions_dir, "ops", "cuda", "*.cu"))
 
-    source_cuda += glob.glob(os.path.join(extensions_dir, 'ops', 'autocast', '*.cpp'))
+    source_cuda += glob.glob(os.path.join(extensions_dir, "ops", "autocast", "*.cpp"))
 
     sources = main_file + source_cpu
     extension = CppExtension
 
-    compile_cpp_tests = os.getenv('WITH_CPP_MODELS_TEST', '0') == '1'
+    compile_cpp_tests = os.getenv("WITH_CPP_MODELS_TEST", "0") == "1"
     if compile_cpp_tests:
-        test_dir = os.path.join(this_dir, 'test')
-        models_dir = os.path.join(this_dir, 'torchvision', 'csrc', 'models')
-        test_file = glob.glob(os.path.join(test_dir, '*.cpp'))
-        source_models = glob.glob(os.path.join(models_dir, '*.cpp'))
+        test_dir = os.path.join(this_dir, "test")
+        models_dir = os.path.join(this_dir, "torchvision", "csrc", "models")
+        test_file = glob.glob(os.path.join(test_dir, "*.cpp"))
+        source_models = glob.glob(os.path.join(models_dir, "*.cpp"))
 
         test_file = [os.path.join(test_dir, s) for s in test_file]
         source_models = [os.path.join(models_dir, s) for s in source_models]
@@ -189,39 +185,38 @@ def get_extensions():
 
     define_macros = []
 
-    extra_compile_args = {'cxx': []}
-    if (torch.cuda.is_available() and ((CUDA_HOME is not None) or is_rocm_pytorch)) \
-            or os.getenv('FORCE_CUDA', '0') == '1':
+    extra_compile_args = {"cxx": []}
+    if (torch.cuda.is_available() and ((CUDA_HOME is not None) or is_rocm_pytorch)) or os.getenv(
+        "FORCE_CUDA", "0"
+    ) == "1":
         extension = CUDAExtension
         sources += source_cuda
         if not is_rocm_pytorch:
-            define_macros += [('WITH_CUDA', None)]
-            nvcc_flags = os.getenv('NVCC_FLAGS', '')
-            if nvcc_flags == '':
+            define_macros += [("WITH_CUDA", None)]
+            nvcc_flags = os.getenv("NVCC_FLAGS", "")
+            if nvcc_flags == "":
                 nvcc_flags = []
             else:
-                nvcc_flags = nvcc_flags.split(' ')
+                nvcc_flags = nvcc_flags.split(" ")
         else:
-            define_macros += [('WITH_HIP', None)]
+            define_macros += [("WITH_HIP", None)]
             nvcc_flags = []
         extra_compile_args["nvcc"] = nvcc_flags
 
-    if sys.platform == 'win32':
-        define_macros += [('torchvision_EXPORTS', None)]
+    if sys.platform == "win32":
+        define_macros += [("torchvision_EXPORTS", None)]
 
-        extra_compile_args['cxx'].append('/MP')
+        extra_compile_args["cxx"].append("/MP")
 
-    debug_mode = os.getenv('DEBUG', '0') == '1'
+    debug_mode = os.getenv("DEBUG", "0") == "1"
     if debug_mode:
         print("Compile in debug mode")
-        extra_compile_args['cxx'].append("-g")
-        extra_compile_args['cxx'].append("-O0")
+        extra_compile_args["cxx"].append("-g")
+        extra_compile_args["cxx"].append("-O0")
         if "nvcc" in extra_compile_args:
             # we have to remove "-OX" and "-g" flag if exists and append
             nvcc_flags = extra_compile_args["nvcc"]
-            extra_compile_args["nvcc"] = [
-                f for f in nvcc_flags if not ("-O" in f or "-g" in f)
-            ]
+            extra_compile_args["nvcc"] = [f for f in nvcc_flags if not ("-O" in f or "-g" in f)]
             extra_compile_args["nvcc"].append("-O0")
             extra_compile_args["nvcc"].append("-g")
 
@@ -231,7 +226,7 @@ def get_extensions():
 
     ext_modules = [
         extension(
-            'torchvision._C',
+            "torchvision._C",
             sorted(sources),
             include_dirs=include_dirs,
             define_macros=define_macros,
@@ -241,7 +236,7 @@ def get_extensions():
     if compile_cpp_tests:
         ext_modules.append(
             extension(
-                'torchvision._C_tests',
+                "torchvision._C_tests",
                 tests,
                 include_dirs=tests_include_dirs,
                 define_macros=define_macros,
@@ -250,12 +245,10 @@ def get_extensions():
         )
 
     # ------------------- Torchvision extra extensions ------------------------
-    vision_include = os.environ.get('TORCHVISION_INCLUDE', None)
-    vision_library = os.environ.get('TORCHVISION_LIBRARY', None)
-    vision_include = (vision_include.split(os.pathsep)
-                      if vision_include is not None else [])
-    vision_library = (vision_library.split(os.pathsep)
-                      if vision_library is not None else [])
+    vision_include = os.environ.get("TORCHVISION_INCLUDE", None)
+    vision_library = os.environ.get("TORCHVISION_LIBRARY", None)
+    vision_include = vision_include.split(os.pathsep) if vision_include is not None else []
+    vision_library = vision_library.split(os.pathsep) if vision_library is not None else []
     include_dirs += vision_include
     library_dirs = vision_library
 
@@ -266,56 +259,49 @@ def get_extensions():
     image_link_flags = []
 
     # Locating libPNG
-    libpng = distutils.spawn.find_executable('libpng-config')
-    pngfix = distutils.spawn.find_executable('pngfix')
+    libpng = distutils.spawn.find_executable("libpng-config")
+    pngfix = distutils.spawn.find_executable("pngfix")
     png_found = libpng is not None or pngfix is not None
-    print('PNG found: {0}'.format(png_found))
+    print("PNG found: {0}".format(png_found))
     if png_found:
         if libpng is not None:
             # Linux / Mac
-            png_version = subprocess.run([libpng, '--version'],
-                                         stdout=subprocess.PIPE)
-            png_version = png_version.stdout.strip().decode('utf-8')
-            print('libpng version: {0}'.format(png_version))
+            png_version = subprocess.run([libpng, "--version"], stdout=subprocess.PIPE)
+            png_version = png_version.stdout.strip().decode("utf-8")
+            print("libpng version: {0}".format(png_version))
             png_version = parse_version(png_version)
             if png_version >= parse_version("1.6.0"):
-                print('Building torchvision with PNG image support')
-                png_lib = subprocess.run([libpng, '--libdir'],
-                                         stdout=subprocess.PIPE)
-                png_lib = png_lib.stdout.strip().decode('utf-8')
-                if 'disabled' not in png_lib:
+                print("Building torchvision with PNG image support")
+                png_lib = subprocess.run([libpng, "--libdir"], stdout=subprocess.PIPE)
+                png_lib = png_lib.stdout.strip().decode("utf-8")
+                if "disabled" not in png_lib:
                     image_library += [png_lib]
-                png_include = subprocess.run([libpng, '--I_opts'],
-                                             stdout=subprocess.PIPE)
-                png_include = png_include.stdout.strip().decode('utf-8')
-                _, png_include = png_include.split('-I')
-                print('libpng include path: {0}'.format(png_include))
+                png_include = subprocess.run([libpng, "--I_opts"], stdout=subprocess.PIPE)
+                png_include = png_include.stdout.strip().decode("utf-8")
+                _, png_include = png_include.split("-I")
+                print("libpng include path: {0}".format(png_include))
                 image_include += [png_include]
-                image_link_flags.append('png')
+                image_link_flags.append("png")
             else:
-                print('libpng installed version is less than 1.6.0, '
-                      'disabling PNG support')
+                print("libpng installed version is less than 1.6.0, " "disabling PNG support")
                 png_found = False
         else:
             # Windows
-            png_lib = os.path.join(
-                os.path.dirname(os.path.dirname(pngfix)), 'lib')
-            png_include = os.path.join(os.path.dirname(
-                os.path.dirname(pngfix)), 'include', 'libpng16')
+            png_lib = os.path.join(os.path.dirname(os.path.dirname(pngfix)), "lib")
+            png_include = os.path.join(os.path.dirname(os.path.dirname(pngfix)), "include", "libpng16")
             image_library += [png_lib]
             image_include += [png_include]
-            image_link_flags.append('libpng')
+            image_link_flags.append("libpng")
 
     # Locating libjpeg
-    (jpeg_found, jpeg_conda,
-     jpeg_include, jpeg_lib) = find_library('jpeglib', vision_include)
+    (jpeg_found, jpeg_conda, jpeg_include, jpeg_lib) = find_library("jpeglib", vision_include)
 
-    print('JPEG found: {0}'.format(jpeg_found))
-    image_macros += [('PNG_FOUND', str(int(png_found)))]
-    image_macros += [('JPEG_FOUND', str(int(jpeg_found)))]
+    print("JPEG found: {0}".format(jpeg_found))
+    image_macros += [("PNG_FOUND", str(int(png_found)))]
+    image_macros += [("JPEG_FOUND", str(int(jpeg_found)))]
     if jpeg_found:
-        print('Building torchvision with JPEG image support')
-        image_link_flags.append('jpeg')
+        print("Building torchvision with JPEG image support")
+        image_link_flags.append("jpeg")
         if jpeg_conda:
             image_library += [jpeg_lib]
             image_include += [jpeg_include]
@@ -323,80 +309,71 @@ def get_extensions():
     # Locating nvjpeg
     # Should be included in CUDA_HOME for CUDA >= 10.1, which is the minimum version we have in the CI
     nvjpeg_found = (
-        extension is CUDAExtension and
-        CUDA_HOME is not None and
-        os.path.exists(os.path.join(CUDA_HOME, 'include', 'nvjpeg.h'))
+        extension is CUDAExtension
+        and CUDA_HOME is not None
+        and os.path.exists(os.path.join(CUDA_HOME, "include", "nvjpeg.h"))
     )
 
-    print('NVJPEG found: {0}'.format(nvjpeg_found))
-    image_macros += [('NVJPEG_FOUND', str(int(nvjpeg_found)))]
+    print("NVJPEG found: {0}".format(nvjpeg_found))
+    image_macros += [("NVJPEG_FOUND", str(int(nvjpeg_found)))]
     if nvjpeg_found:
-        print('Building torchvision with NVJPEG image support')
-        image_link_flags.append('nvjpeg')
-
-    image_path = os.path.join(extensions_dir, 'io', 'image')
-    image_src = (glob.glob(os.path.join(image_path, '*.cpp')) + glob.glob(os.path.join(image_path, 'cpu', '*.cpp'))
-                 + glob.glob(os.path.join(image_path, 'cuda', '*.cpp')))
+        print("Building torchvision with NVJPEG image support")
+        image_link_flags.append("nvjpeg")
+
+    image_path = os.path.join(extensions_dir, "io", "image")
+    image_src = (
+        glob.glob(os.path.join(image_path, "*.cpp"))
+        + glob.glob(os.path.join(image_path, "cpu", "*.cpp"))
+        + glob.glob(os.path.join(image_path, "cuda", "*.cpp"))
+    )
 
     if png_found or jpeg_found:
-        ext_modules.append(extension(
-            'torchvision.image',
-            image_src,
-            include_dirs=image_include + include_dirs + [image_path],
-            library_dirs=image_library + library_dirs,
-            define_macros=image_macros,
-            libraries=image_link_flags,
-            extra_compile_args=extra_compile_args
-        ))
-
-    ffmpeg_exe = distutils.spawn.find_executable('ffmpeg')
+        ext_modules.append(
+            extension(
+                "torchvision.image",
+                image_src,
+                include_dirs=image_include + include_dirs + [image_path],
+                library_dirs=image_library + library_dirs,
+                define_macros=image_macros,
+                libraries=image_link_flags,
+                extra_compile_args=extra_compile_args,
+            )
+        )
+
+    ffmpeg_exe = distutils.spawn.find_executable("ffmpeg")
     has_ffmpeg = ffmpeg_exe is not None
     # FIXME: Building torchvision with ffmpeg on MacOS or with Python 3.9
     # FIXME: causes crash. See the following GitHub issues for more details.
     # FIXME: https://github.com/pytorch/pytorch/issues/65000
     # FIXME: https://github.com/pytorch/vision/issues/3367
-    if sys.platform != 'linux' or (
-            sys.version_info.major == 3 and sys.version_info.minor == 9):
+    if sys.platform != "linux" or (sys.version_info.major == 3 and sys.version_info.minor == 9):
         has_ffmpeg = False
     if has_ffmpeg:
         try:
             # This is to check if ffmpeg is installed properly.
             subprocess.check_output(["ffmpeg", "-version"])
         except subprocess.CalledProcessError:
-            print('Error fetching ffmpeg version, ignoring ffmpeg.')
+            print("Error fetching ffmpeg version, ignoring ffmpeg.")
             has_ffmpeg = False
 
     print("FFmpeg found: {}".format(has_ffmpeg))
 
     if has_ffmpeg:
-        ffmpeg_libraries = {
-            'libavcodec',
-            'libavformat',
-            'libavutil',
-            'libswresample',
-            'libswscale'
-        }
+        ffmpeg_libraries = {"libavcodec", "libavformat", "libavutil", "libswresample", "libswscale"}
 
         ffmpeg_bin = os.path.dirname(ffmpeg_exe)
         ffmpeg_root = os.path.dirname(ffmpeg_bin)
-        ffmpeg_include_dir = os.path.join(ffmpeg_root, 'include')
-        ffmpeg_library_dir = os.path.join(ffmpeg_root, 'lib')
+        ffmpeg_include_dir = os.path.join(ffmpeg_root, "include")
+        ffmpeg_library_dir = os.path.join(ffmpeg_root, "lib")
 
-        gcc = distutils.spawn.find_executable('gcc')
-        platform_tag = subprocess.run(
-            [gcc, '-print-multiarch'], stdout=subprocess.PIPE)
-        platform_tag = platform_tag.stdout.strip().decode('utf-8')
+        gcc = distutils.spawn.find_executable("gcc")
+        platform_tag = subprocess.run([gcc, "-print-multiarch"], stdout=subprocess.PIPE)
+        platform_tag = platform_tag.stdout.strip().decode("utf-8")
 
         if platform_tag:
             # Most probably a Debian-based distribution
-            ffmpeg_include_dir = [
-                ffmpeg_include_dir,
-                os.path.join(ffmpeg_include_dir, platform_tag)
-            ]
-            ffmpeg_library_dir = [
-                ffmpeg_library_dir,
-                os.path.join(ffmpeg_library_dir, platform_tag)
-            ]
+            ffmpeg_include_dir = [ffmpeg_include_dir, os.path.join(ffmpeg_include_dir, platform_tag)]
+            ffmpeg_library_dir = [ffmpeg_library_dir, os.path.join(ffmpeg_library_dir, platform_tag)]
         else:
             ffmpeg_include_dir = [ffmpeg_include_dir]
             ffmpeg_library_dir = [ffmpeg_library_dir]
@@ -405,11 +382,11 @@ def get_extensions():
         for library in ffmpeg_libraries:
             library_found = False
             for search_path in ffmpeg_include_dir + include_dirs:
-                full_path = os.path.join(search_path, library, '*.h')
+                full_path = os.path.join(search_path, library, "*.h")
                 library_found |= len(glob.glob(full_path)) > 0
 
             if not library_found:
-                print(f'{library} header files were not found, disabling ffmpeg support')
+                print(f"{library} header files were not found, disabling ffmpeg support")
                 has_ffmpeg = False
 
     if has_ffmpeg:
@@ -417,22 +394,21 @@ def get_extensions():
         print("ffmpeg library_dir: {}".format(ffmpeg_library_dir))
 
         # TorchVision base decoder + video reader
-        video_reader_src_dir = os.path.join(this_dir, 'torchvision', 'csrc', 'io', 'video_reader')
+        video_reader_src_dir = os.path.join(this_dir, "torchvision", "csrc", "io", "video_reader")
         video_reader_src = glob.glob(os.path.join(video_reader_src_dir, "*.cpp"))
-        base_decoder_src_dir = os.path.join(this_dir, 'torchvision', 'csrc', 'io', 'decoder')
-        base_decoder_src = glob.glob(
-            os.path.join(base_decoder_src_dir, "*.cpp"))
+        base_decoder_src_dir = os.path.join(this_dir, "torchvision", "csrc", "io", "decoder")
+        base_decoder_src = glob.glob(os.path.join(base_decoder_src_dir, "*.cpp"))
         # Torchvision video API
-        videoapi_src_dir = os.path.join(this_dir, 'torchvision', 'csrc', 'io', 'video')
+        videoapi_src_dir = os.path.join(this_dir, "torchvision", "csrc", "io", "video")
         videoapi_src = glob.glob(os.path.join(videoapi_src_dir, "*.cpp"))
         # exclude tests
-        base_decoder_src = [x for x in base_decoder_src if '_test.cpp' not in x]
+        base_decoder_src = [x for x in base_decoder_src if "_test.cpp" not in x]
 
         combined_src = video_reader_src + base_decoder_src + videoapi_src
 
         ext_modules.append(
             CppExtension(
-                'torchvision.video_reader',
+                "torchvision.video_reader",
                 combined_src,
                 include_dirs=[
                     base_decoder_src_dir,
@@ -440,18 +416,18 @@ def get_extensions():
                     videoapi_src_dir,
                     extensions_dir,
                     *ffmpeg_include_dir,
-                    *include_dirs
+                    *include_dirs,
                 ],
                 library_dirs=ffmpeg_library_dir + library_dirs,
                 libraries=[
-                    'avcodec',
-                    'avformat',
-                    'avutil',
-                    'swresample',
-                    'swscale',
+                    "avcodec",
+                    "avformat",
+                    "avutil",
+                    "swresample",
+                    "swscale",
                 ],
-                extra_compile_args=["-std=c++14"] if os.name != 'nt' else ['/std:c++14', '/MP'],
-                extra_link_args=["-std=c++14" if os.name != 'nt' else '/std:c++14'],
+                extra_compile_args=["-std=c++14"] if os.name != "nt" else ["/std:c++14", "/MP"],
+                extra_link_args=["-std=c++14" if os.name != "nt" else "/std:c++14"],
             )
         )
 
@@ -460,9 +436,9 @@ def get_extensions():
 
 class clean(distutils.command.clean.clean):
     def run(self):
-        with open('.gitignore', 'r') as f:
+        with open(".gitignore", "r") as f:
             ignores = f.read()
-            for wildcard in filter(None, ignores.split('\n')):
+            for wildcard in filter(None, ignores.split("\n")):
                 for filename in glob.glob(wildcard):
                     try:
                         os.remove(filename)
@@ -478,25 +454,22 @@ if __name__ == "__main__":
 
     write_version_file()
 
-    with open('README.rst') as f:
+    with open("README.rst") as f:
         readme = f.read()
 
     setup(
         # Metadata
         name=package_name,
         version=version,
-        author='PyTorch Core Team',
-        author_email='soumith@pytorch.org',
-        url='https://github.com/pytorch/vision',
-        description='image and video datasets and models for torch deep learning',
+        author="PyTorch Core Team",
+        author_email="soumith@pytorch.org",
+        url="https://github.com/pytorch/vision",
+        description="image and video datasets and models for torch deep learning",
         long_description=readme,
-        license='BSD',
-
+        license="BSD",
         # Package info
-        packages=find_packages(exclude=('test',)),
-        package_data={
-            package_name: ['*.dll', '*.dylib', '*.so', '*.categories']
-        },
+        packages=find_packages(exclude=("test",)),
+        package_data={package_name: ["*.dll", "*.dylib", "*.so", "*.categories"]},
         zip_safe=False,
         install_requires=requirements,
         extras_require={
@@ -504,7 +477,7 @@ if __name__ == "__main__":
         },
         ext_modules=get_extensions(),
         cmdclass={
-            'build_ext': BuildExtension.with_options(no_python_abi_suffix=True),
-            'clean': clean,
-        }
+            "build_ext": BuildExtension.with_options(no_python_abi_suffix=True),
+            "clean": clean,
+        },
     )

test/common_utils.py

+import argparse
+import contextlib
+import functools
+import inspect
 import os
+import random
 import shutil
+import sys
 import tempfile
-import contextlib
 import unittest
+from collections import OrderedDict
+from numbers import Number
+
+import numpy as np
 import pytest
-import argparse
-import sys
 import torch
-import __main__
-import random
-import inspect
-import functools
-
-from numbers import Number
+from PIL import Image
 from torch._six import string_classes
-from collections import OrderedDict
 from torchvision import io
 
-import numpy as np
-from PIL import Image
+import __main__
 
 
-IN_CIRCLE_CI = os.getenv("CIRCLECI", False) == 'true'
+IN_CIRCLE_CI = os.getenv("CIRCLECI", False) == "true"
 IN_RE_WORKER = os.environ.get("INSIDE_RE_WORKER") is not None
 IN_FBCODE = os.environ.get("IN_FBCODE_TORCHVISION") == "1"
-CUDA_NOT_AVAILABLE_MSG = 'CUDA device not available'
+CUDA_NOT_AVAILABLE_MSG = "CUDA device not available"
 CIRCLECI_GPU_NO_CUDA_MSG = "We're in a CircleCI GPU machine, and this test doesn't need cuda."
 
 
@@ -95,7 +95,7 @@ def freeze_rng_state():
 
 def cycle_over(objs):
     for idx, obj1 in enumerate(objs):
-        for obj2 in objs[:idx] + objs[idx + 1:]:
+        for obj2 in objs[:idx] + objs[idx + 1 :]:
             yield obj1, obj2
 
 
@@ -117,11 +117,13 @@ def disable_console_output():
 
 def cpu_and_gpu():
     import pytest  # noqa
-    return ('cpu', pytest.param('cuda', marks=pytest.mark.needs_cuda))
+
+    return ("cpu", pytest.param("cuda", marks=pytest.mark.needs_cuda))
 
 
 def needs_cuda(test_func):
     import pytest  # noqa
+
     return pytest.mark.needs_cuda(test_func)
 
 
@@ -139,12 +141,7 @@ def _create_data(height=3, width=3, channels=3, device="cpu"):
 
 def _create_data_batch(height=3, width=3, channels=3, num_samples=4, device="cpu"):
     # TODO: When all relevant tests are ported to pytest, turn this into a module-level fixture
-    batch_tensor = torch.randint(
-        0, 256,
-        (num_samples, channels, height, width),
-        dtype=torch.uint8,
-        device=device
-    )
+    batch_tensor = torch.randint(0, 256, (num_samples, channels, height, width), dtype=torch.uint8, device=device)
     return batch_tensor
 
 
@@ -180,8 +177,9 @@ def _assert_equal_tensor_to_pil(tensor, pil_image, msg=None):
     assert_equal(tensor.cpu(), pil_tensor, msg=msg)
 
 
-def _assert_approx_equal_tensor_to_pil(tensor, pil_image, tol=1e-5, msg=None, agg_method="mean",
-                                       allowed_percentage_diff=None):
+def _assert_approx_equal_tensor_to_pil(
+    tensor, pil_image, tol=1e-5, msg=None, agg_method="mean", allowed_percentage_diff=None
+):
     # TODO: we could just merge this into _assert_equal_tensor_to_pil
     np_pil_image = np.array(pil_image)
     if np_pil_image.ndim == 2: