[FX] remove op_lowering_disallow_list and format revert (#1261)

* sync to fb master

* reverse _compile.py change

* comment line length to use default

* update nightly pytorch to 0810

* black formatting

* update

.circleci/config.yml

@@ -751,7 +751,7 @@ parameters:
   # Nightly platform config
   torch-nightly-build:
     type: string
-    default: "1.13.0.dev20220731+cu113"
+    default: "1.13.0.dev20220810+cu113"
   torch-nightly-build-index:
     type: string
     default: "https://download.pytorch.org/whl/nightly/cu113"

docsrc/conf.py

@@ -99,7 +99,9 @@
 }
 
 html_show_sourcelink = True
-html_sidebars = {"**": ["logo-text.html", "globaltoc.html", "localtoc.html", "searchbox.html"]}
+html_sidebars = {
+    "**": ["logo-text.html", "globaltoc.html", "localtoc.html", "searchbox.html"]
+}
 
 # extensions.append("sphinx_material")
 html_theme_path = [pytorch_sphinx_theme.get_html_theme_path()]
@@ -183,7 +185,15 @@ def handle_item(fieldarg, content):
                 typename = typename.replace("long", "python:long")
                 typename = typename.replace("float", "python:float")
                 typename = typename.replace("type", "python:type")
-                par.extend(self.make_xrefs(self.typerolename, domain, typename, addnodes.literal_emphasis, **kw))
+                par.extend(
+                    self.make_xrefs(
+                        self.typerolename,
+                        domain,
+                        typename,
+                        addnodes.literal_emphasis,
+                        **kw
+                    )
+                )
             else:
                 par += fieldtype
             par += nodes.Text(")")

examples/custom_converters/elu_model.py

@@ -2,7 +2,9 @@
 import torch_tensorrt
 
 # After "python3 setup install", you should find this .so file under generated "build" directory
-torch.ops.load_library("./elu_converter/build/lib.linux-x86_64-3.6/elu_converter.cpython-36m-x86_64-linux-gnu.so")
+torch.ops.load_library(
+    "./elu_converter/build/lib.linux-x86_64-3.6/elu_converter.cpython-36m-x86_64-linux-gnu.so"
+)
 
 
 class Elu(torch.nn.Module):

examples/fx/fx2trt_example.py

@@ -141,4 +141,6 @@ def get_input(self, inputs):
 
 # Make sure the results match
 regular_model_output = model(*inputs)
-torch.testing.assert_close(reload_model_output, regular_model_output, atol=3e-3, rtol=1e-2)
+torch.testing.assert_close(
+    reload_model_output, regular_model_output, atol=3e-3, rtol=1e-2
+)

examples/fx/hugging_face_torchdynamo_example.py

@@ -353,14 +353,18 @@ def run_all_eval(args, optimize_ctx, optimize_name, dtype):
             eval_inputs = (input_ids,)
 
         # Correctness check
-        is_accurate = check_correctness(args, model, eval_inputs, optimize_ctx, optimize_name)
+        is_accurate = check_correctness(
+            args, model, eval_inputs, optimize_ctx, optimize_name
+        )
         # Profile eager
         t, m = bench_model_eval(args, "eager", model, eval_inputs, NullContext())
         results.append(create_record(model_name, dtype, is_accurate, "eager", t, m))
 
         # Profile Dynamo nvfuser
         t, m = bench_model_eval(args, optimize_name, model, eval_inputs, optimize_ctx)
-        results.append(create_record(model_name, dtype, is_accurate, optimize_name, t, m))
+        results.append(
+            create_record(model_name, dtype, is_accurate, optimize_name, t, m)
+        )
 
         # calculate relative improvements
         base_r = results[-2]
@@ -412,7 +416,9 @@ def main():
     if optimize_name == "dynamo_fx2trt_fp32":
         experiment = partial(experiment, dtype=torch.float32)
 
-    experiment = partial(experiment, optimize_ctx=optimize_ctx, optimize_name=optimize_name)
+    experiment = partial(
+        experiment, optimize_ctx=optimize_ctx, optimize_name=optimize_name
+    )
     experiment(args)
 
 

examples/fx/lower_example.py

@@ -125,7 +125,10 @@ def benchmark(
         ),
     ]
 
-    results = [run_configuration_benchmark(deepcopy(model), inputs, conf_) for conf_ in configurations]
+    results = [
+        run_configuration_benchmark(deepcopy(model), inputs, conf_)
+        for conf_ in configurations
+    ]
 
     for res in results:
         print(res.format())

examples/fx/quantized_resnet_test.py

@@ -108,7 +108,9 @@ def build_int8_trt_implicit_quant(rn18):
         InputTensorSpec.from_tensors([data]),
         logger_level=trt.Logger.VERBOSE,
     )
-    interpreter_result = interp.run(lower_precision=LowerPrecision.INT8, strict_type_constraints=True)
+    interpreter_result = interp.run(
+        lower_precision=LowerPrecision.INT8, strict_type_constraints=True
+    )
     trt_mod = TRTModule(
         interpreter_result.engine,
         interpreter_result.input_names,

examples/fx/torch_trt_simple_example.py

@@ -12,35 +12,63 @@ def test_torch_tensorrt(model, inputs):
     # fp32 test
     with torch.inference_mode():
         ref_fp32 = model_ts(*inputs_ts)
-    trt_ts_module = torch_tensorrt.compile(model_ts, inputs=inputs_ts, enabled_precisions={torch.float32})
+    trt_ts_module = torch_tensorrt.compile(
+        model_ts, inputs=inputs_ts, enabled_precisions={torch.float32}
+    )
     result_fp32 = trt_ts_module(*inputs_ts)
-    assert torch.nn.functional.cosine_similarity(ref_fp32.flatten(), result_fp32.flatten(), dim=0) > 0.9999
+    assert (
+        torch.nn.functional.cosine_similarity(
+            ref_fp32.flatten(), result_fp32.flatten(), dim=0
+        )
+        > 0.9999
+    )
     # fp16 test
     model_ts = model_ts.half()
     inputs_ts = [i.cuda().half() for i in inputs_ts]
     with torch.inference_mode():
         ref_fp16 = model_ts(*inputs_ts)
-    trt_ts_module = torch_tensorrt.compile(model_ts, inputs=inputs_ts, enabled_precisions={torch.float16})
+    trt_ts_module = torch_tensorrt.compile(
+        model_ts, inputs=inputs_ts, enabled_precisions={torch.float16}
+    )
     result_fp16 = trt_ts_module(*inputs_ts)
-    assert torch.nn.functional.cosine_similarity(ref_fp16.flatten(), result_fp16.flatten(), dim=0) > 0.99
+    assert (
+        torch.nn.functional.cosine_similarity(
+            ref_fp16.flatten(), result_fp16.flatten(), dim=0
+        )
+        > 0.99
+    )
 
     # FX path
     model_fx = copy.deepcopy(model)
     inputs_fx = copy.deepcopy(inputs)
     # fp32 test
     with torch.inference_mode():
         ref_fp32 = model_fx(*inputs_fx)
-    trt_fx_module = torch_tensorrt.compile(model_fx, ir="fx", inputs=inputs_fx, enabled_precisions={torch.float32})
+    trt_fx_module = torch_tensorrt.compile(
+        model_fx, ir="fx", inputs=inputs_fx, enabled_precisions={torch.float32}
+    )
     result_fp32 = trt_fx_module(*inputs_fx)
-    assert torch.nn.functional.cosine_similarity(ref_fp32.flatten(), result_fp32.flatten(), dim=0) > 0.9999
+    assert (
+        torch.nn.functional.cosine_similarity(
+            ref_fp32.flatten(), result_fp32.flatten(), dim=0
+        )
+        > 0.9999
+    )
     # fp16 test
     model_fx = model_fx.cuda().half()
     inputs_fx = [i.cuda().half() for i in inputs_fx]
     with torch.inference_mode():
         ref_fp16 = model_fx(*inputs_fx)
-    trt_fx_module = torch_tensorrt.compile(model_fx, ir="fx", inputs=inputs_fx, enabled_precisions={torch.float16})
+    trt_fx_module = torch_tensorrt.compile(
+        model_fx, ir="fx", inputs=inputs_fx, enabled_precisions={torch.float16}
+    )
     result_fp16 = trt_fx_module(*inputs_fx)
-    assert torch.nn.functional.cosine_similarity(ref_fp16.flatten(), result_fp16.flatten(), dim=0) > 0.99
+    assert (
+        torch.nn.functional.cosine_similarity(
+            ref_fp16.flatten(), result_fp16.flatten(), dim=0
+        )
+        > 0.99
+    )
 
 
 if __name__ == "__main__":

examples/fx/torchdynamo_example.py

@@ -142,7 +142,10 @@ def benchmark(
         ),
     ]
 
-    results = [run_configuration_benchmark(deepcopy(model), inputs, conf_) for conf_ in configurations]
+    results = [
+        run_configuration_benchmark(deepcopy(model), inputs, conf_)
+        for conf_ in configurations
+    ]
 
     for res in results:
         print(res.format())

examples/int8/training/vgg16/export_ckpt.py

@@ -75,7 +75,9 @@ def test(model, dataloader, crit):
     ),
 )
 
-testing_dataloader = torch.utils.data.DataLoader(testing_dataset, batch_size=32, shuffle=False, num_workers=2)
+testing_dataloader = torch.utils.data.DataLoader(
+    testing_dataset, batch_size=32, shuffle=False, num_workers=2
+)
 
 crit = torch.nn.CrossEntropyLoss()
 

examples/int8/training/vgg16/export_qat.py

@@ -72,7 +72,9 @@ def test(model, dataloader, crit):
     ),
 )
 
-testing_dataloader = torch.utils.data.DataLoader(testing_dataset, batch_size=32, shuffle=False, num_workers=2)
+testing_dataloader = torch.utils.data.DataLoader(
+    testing_dataset, batch_size=32, shuffle=False, num_workers=2
+)
 
 crit = torch.nn.CrossEntropyLoss()
 

examples/int8/training/vgg16/finetune_qat.py

@@ -21,14 +21,20 @@
 
 from vgg16 import vgg16
 
-PARSER = argparse.ArgumentParser(description="VGG16 example to use with Torch-TensorRT PTQ")
-PARSER.add_argument("--epochs", default=100, type=int, help="Number of total epochs to train")
+PARSER = argparse.ArgumentParser(
+    description="VGG16 example to use with Torch-TensorRT PTQ"
+)
+PARSER.add_argument(
+    "--epochs", default=100, type=int, help="Number of total epochs to train"
+)
 PARSER.add_argument(
     "--enable_qat",
     action="store_true",
     help="Enable quantization aware training. This is recommended to perform on a pre-trained model.",
 )
-PARSER.add_argument("--batch-size", default=128, type=int, help="Batch size to use when training")
+PARSER.add_argument(
+    "--batch-size", default=128, type=int, help="Batch size to use when training"
+)
 PARSER.add_argument("--lr", default=0.1, type=float, help="Initial learning rate")
 PARSER.add_argument("--drop-ratio", default=0.0, type=float, help="Dropout ratio")
 PARSER.add_argument("--momentum", default=0.9, type=float, help="Momentum")
@@ -194,7 +200,9 @@ def main():
                 transforms.RandomCrop(32, padding=4),
                 transforms.RandomHorizontalFlip(),
                 transforms.ToTensor(),
-                transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
+                transforms.Normalize(
+                    (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
+                ),
             ]
         ),
     )
@@ -209,7 +217,9 @@ def main():
         transform=transforms.Compose(
             [
                 transforms.ToTensor(),
-                transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
+                transforms.Normalize(
+                    (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
+                ),
             ]
         ),
     )
@@ -309,9 +319,14 @@ def train(model, dataloader, crit, opt, epoch):
 
         running_loss += loss.item()
         if batch % 50 == 49:
-            writer.add_scalar("Training Loss", running_loss / 100, epoch * len(dataloader) + batch)
+            writer.add_scalar(
+                "Training Loss", running_loss / 100, epoch * len(dataloader) + batch
+            )
             writer.close()
-            print("Batch: [%5d | %5d] loss: %.3f" % (batch + 1, len(dataloader), running_loss / 100))
+            print(
+                "Batch: [%5d | %5d] loss: %.3f"
+                % (batch + 1, len(dataloader), running_loss / 100)
+            )
             running_loss = 0.0
 
 

examples/int8/training/vgg16/main.py

@@ -15,9 +15,15 @@
 
 from vgg16 import vgg16
 
-PARSER = argparse.ArgumentParser(description="VGG16 example to use with Torch-TensorRT PTQ")
-PARSER.add_argument("--epochs", default=100, type=int, help="Number of total epochs to train")
-PARSER.add_argument("--batch-size", default=128, type=int, help="Batch size to use when training")
+PARSER = argparse.ArgumentParser(
+    description="VGG16 example to use with Torch-TensorRT PTQ"
+)
+PARSER.add_argument(
+    "--epochs", default=100, type=int, help="Number of total epochs to train"
+)
+PARSER.add_argument(
+    "--batch-size", default=128, type=int, help="Batch size to use when training"
+)
 PARSER.add_argument("--lr", default=0.1, type=float, help="Initial learning rate")
 PARSER.add_argument("--drop-ratio", default=0.0, type=float, help="Dropout ratio")
 PARSER.add_argument("--momentum", default=0.9, type=float, help="Momentum")
@@ -89,7 +95,9 @@ def main():
                 transforms.RandomCrop(32, padding=4),
                 transforms.RandomHorizontalFlip(),
                 transforms.ToTensor(),
-                transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
+                transforms.Normalize(
+                    (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
+                ),
             ]
         ),
     )
@@ -104,7 +112,9 @@ def main():
         transform=transforms.Compose(
             [
                 transforms.ToTensor(),
-                transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
+                transforms.Normalize(
+                    (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
+                ),
             ]
         ),
     )
@@ -182,9 +192,14 @@ def train(model, dataloader, crit, opt, epoch):
 
         running_loss += loss.item()
         if batch % 50 == 49:
-            writer.add_scalar("Training Loss", running_loss / 100, epoch * len(dataloader) + batch)
+            writer.add_scalar(
+                "Training Loss", running_loss / 100, epoch * len(dataloader) + batch
+            )
             writer.close()
-            print("Batch: [%5d | %5d] loss: %.3f" % (batch + 1, len(dataloader), running_loss / 100))
+            print(
+                "Batch: [%5d | %5d] loss: %.3f"
+                % (batch + 1, len(dataloader), running_loss / 100)
+            )
             running_loss = 0.0
 
 

examples/int8/training/vgg16/test_qat.py

@@ -71,7 +71,9 @@ def test(model, dataloader, crit):
     ),
 )
 
-testing_dataloader = torch.utils.data.DataLoader(testing_dataset, batch_size=32, shuffle=False, num_workers=2)
+testing_dataloader = torch.utils.data.DataLoader(
+    testing_dataset, batch_size=32, shuffle=False, num_workers=2
+)
 
 crit = torch.nn.CrossEntropyLoss()
 
@@ -94,6 +96,8 @@ def test(model, dataloader, crit):
 }
 new_mod = torch.jit.load("trained_vgg16_qat.jit.pt")
 trt_ts_module = torchtrt.compile(new_mod, **compile_settings)
-testing_dataloader = torch.utils.data.DataLoader(testing_dataset, batch_size=1, shuffle=False, num_workers=2)
+testing_dataloader = torch.utils.data.DataLoader(
+    testing_dataset, batch_size=1, shuffle=False, num_workers=2
+)
 test_loss, test_acc = test(trt_ts_module, testing_dataloader, crit)
 print("[TRTorch] Test Loss: {:.5f} Test Acc: {:.2f}%".format(test_loss, 100 * test_acc))