Repository: pytorch/TensorRT

Issue:
Support TensorRT Post Training Quantization
We want to be able to use TensorRT's PTQ using a PyTorch dataset to support INT8 execution. 
FP16 Execution produces non accurate results
Right now the same workflow to run FP32 models when used in FP16 gives bad results.
FP16 Execution produces non accurate results
Right now the same workflow to run FP32 models when used in FP16 gives bad results.


Discussion:

Resolving patch:
diff --git a/cpp/ptq/training/vgg16/export_ckpt.py b/cpp/ptq/training/vgg16/export_ckpt.py
new file mode 100644
--- /dev/null
+++ b/cpp/ptq/training/vgg16/export_ckpt.py
@@ -0,0 +1,76 @@
+import argparse
+import os
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.data as data
+import torchvision.transforms as transforms
+import torchvision.datasets as datasets
+
+from vgg16 import vgg16
+
+def test(model, dataloader, crit):
+    global writer
+    global classes
+    total = 0
+    correct = 0
+    loss = 0.0
+    class_probs = []
+    class_preds = []
+
+    with torch.no_grad():
+        for data, labels in dataloader:
+            data, labels = data.cuda(), labels.cuda(async=True)
+            out = model(data)
+            loss += crit(out, labels)
+            preds = torch.max(out, 1)[1]
+            class_probs.append([F.softmax(i, dim=0) for i in out])
+            class_preds.append(preds)
+            total += labels.size(0)
+            correct += (preds == labels).sum().item()
+
+    test_probs = torch.cat([torch.stack(batch) for batch in class_probs])
+    test_preds = torch.cat(class_preds)
+    return loss / total, correct / total
+
+PARSER = argparse.ArgumentParser(description="Export trained VGG")
+PARSER.add_argument('ckpt', type=str, help="Path to saved checkpoint")
+
+args = PARSER.parse_args()
+model = vgg16(num_classes=10, init_weights=False)
+model = model.cuda()
+
+ckpt = torch.load(args.ckpt)
+weights = ckpt["model_state_dict"]
+
+if torch.cuda.device_count() > 1:
+    from collections import OrderedDict
+    new_state_dict = OrderedDict()
+    for k, v in weights.items():
+        name = k[7:] # remove `module.`
+        new_state_dict[name] = v
+    weights = new_state_dict
+
+model.load_state_dict(weights)
+
+# Setting eval here causes both JIT and TRT accuracy to tank in LibTorch will follow up with PyTorch Team
+#model.eval()
+
+jit_model = torch.jit.trace(model, torch.rand([32, 3, 32, 32]).to("cuda"))
+jit_model.eval()
+
+testing_dataset = datasets.CIFAR10(root='./data', train=False, download=True,
+                                    transform=transforms.Compose([
+                                    transforms.ToTensor(),
+                                    transforms.Normalize((0.4914, 0.4822, 0.4465),
+                                                         (0.2023, 0.1994, 0.2010))]))
+
+testing_dataloader = torch.utils.data.DataLoader(testing_dataset, batch_size=32,
+                                                 shuffle=False, num_workers=2)
+
+crit = torch.nn.CrossEntropyLoss()
+
+test_loss, test_acc = test(jit_model, testing_dataloader, crit)
+print("[JIT] Test Loss: {:.5f} Test Acc: {:.2f}%".format(test_loss, 100 * test_acc))
+torch.jit.save(jit_model, "trained_vgg16.jit.pt")
diff --git a/cpp/ptq/training/vgg16/main.py b/cpp/ptq/training/vgg16/main.py
new file mode 100644
--- /dev/null
+++ b/cpp/ptq/training/vgg16/main.py
@@ -0,0 +1,208 @@
+import argparse
+import os
+import random
+from datetime import datetime
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import torch.utils.data as data
+import torchvision.transforms as transforms
+import torchvision.datasets as datasets
+
+from torch.utils.tensorboard import SummaryWriter
+
+from vgg16 import vgg16
+
+PARSER = argparse.ArgumentParser(description="VGG16 example to use with TRTorch 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., type=float, help="Dropout ratio")
+PARSER.add_argument('--momentum', default=0.9, type=float, help="Momentum")
+PARSER.add_argument('--weight-decay', default=5e-4, type=float, help="Weight decay")
+PARSER.add_argument('--ckpt-dir', default="/tmp/vgg16_ckpts", type=str, help="Path to save checkpoints (saved every 10 epochs)")
+PARSER.add_argument('--start-from', default=0, type=int, help="Epoch to resume from (requires a checkpoin in the providied checkpoi")
+PARSER.add_argument('--seed', type=int, help='Seed value for rng')
+PARSER.add_argument('--tensorboard', type=str, default='/tmp/vgg16_logs', help='Location for tensorboard info')
+
+args = PARSER.parse_args()
+for arg in vars(args):
+    print(' {} {}'.format(arg, getattr(args, arg)))
+state = {k: v for k, v in args._get_kwargs()}
+
+if args.seed is None:
+    args.seed = random.randint(1, 10000)
+random.seed(args.seed)
+torch.manual_seed(args.seed)
+torch.cuda.manual_seed_all(args.seed)
+print("RNG seed used: ", args.seed)
+
+now = datetime.now()
+
+timestamp = datetime.timestamp(now)
+
+writer = SummaryWriter(args.tensorboard + '/test_' + str(timestamp))
+classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
+
+
+def main():
+    global state
+    global classes
+    global writer
+    if not os.path.isdir(args.ckpt_dir):
+        os.makedirs(args.ckpt_dir)
+
+    training_dataset = datasets.CIFAR10(root='./data', train=True,
+                                        download=True, transform=transforms.Compose([
+                                            transforms.RandomCrop(32, padding=4),
+                                            transforms.RandomHorizontalFlip(),
+                                            transforms.ToTensor(),
+                                            transforms.Normalize((0.4914, 0.4822, 0.4465),
+                                                                 (0.2023, 0.1994, 0.2010)),
+                                        ]))
+    training_dataloader = torch.utils.data.DataLoader(training_dataset, batch_size=args.batch_size,
+                                                      shuffle=True, num_workers=2)
+
+    testing_dataset = datasets.CIFAR10(root='./data', train=False, download=True,
+                                       transform=transforms.Compose([
+                                            transforms.ToTensor(),
+                                            transforms.Normalize((0.4914, 0.4822, 0.4465),
+                                                                 (0.2023, 0.1994, 0.2010)),
+                                        ]))
+
+    testing_dataloader = torch.utils.data.DataLoader(testing_dataset, batch_size=args.batch_size,
+                                                     shuffle=False, num_workers=2)
+
+    num_classes = len(classes)
+
+    model = vgg16(num_classes=num_classes, init_weights=False)
+    model = model.cuda()
+
+    data = iter(training_dataloader)
+    images, _ = data.next()
+
+    writer.add_graph(model, images.cuda())
+    writer.close()
+
+    crit = nn.CrossEntropyLoss()
+    opt = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
+
+    if torch.cuda.device_count() > 1:
+        model = nn.DataParallel(model)
+
+    if args.start_from != 0:
+        ckpt_file = args.ckpt_dir + '/ckpt_epoch' + str(args.start_from) + '.pth'
+        print('Loading from checkpoint {}'.format(ckpt_file))
+        assert(os.path.isfile(ckpt_file))
+        ckpt = torch.load(ckpt_file)
+        model.load_state_dict(ckpt["model_state_dict"])
+        opt.load_state_dict(ckpt["opt_state_dict"])
+        state = ckpt["state"]
+
+    for epoch in range(args.start_from, args.epochs):
+        adjust_lr(opt, epoch)
+        writer.add_scalar('Learning Rate', state["lr"], epoch)
+        writer.close()
+        print('Epoch: [%5d / %5d] LR: %f' % (epoch + 1, args.epochs, state['lr']))
+
+        train(model, training_dataloader, crit, opt, epoch)
+        test_loss, test_acc = test(model, testing_dataloader, crit, epoch)
+
+        print("Test Loss: {:.5f} Test Acc: {:.2f}%".format(test_loss, 100 * test_acc))
+
+        if epoch % 10 == 9:
+            save_checkpoint({
+                'epoch': epoch + 1,
+                'model_state_dict': model.state_dict(),
+                'acc': test_acc,
+                'opt_state_dict' : opt.state_dict(),
+                'state': state
+            }, ckpt_dir=args.ckpt_dir)
+
+def train(model, dataloader, crit, opt, epoch):
+    global writer
+    model.train()
+    running_loss = 0.0
+    for batch, (data, labels) in enumerate(dataloader):
+        data, labels = data.cuda(), labels.cuda(async=True)
+        opt.zero_grad()
+        out = model(data)
+        loss = crit(out, labels)
+        loss.backward()
+        opt.step()
+
+        running_loss += loss.item()
+        if batch % 50 == 49:
+            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))
+            running_loss = 0.0
+
+def test(model, dataloader, crit, epoch):
+    global writer
+    global classes
+    total = 0
+    correct = 0
+    loss = 0.0
+    class_probs = []
+    class_preds = []
+    model.eval()
+    with torch.no_grad():
+        for data, labels in dataloader:
+            data, labels = data.cuda(), labels.cuda(async=True)
+            out = model(data)
+            loss += crit(out, labels)
+            preds = torch.max(out, 1)[1]
+            class_probs.append([F.softmax(i, dim=0) for i in out])
+            class_preds.append(preds)
+            total += labels.size(0)
+            correct += (preds == labels).sum().item()
+
+    writer.add_scalar('Testing Loss', loss / total, epoch)
+    writer.close()
+
+    writer.add_scalar('Testing Accuracy', correct / total * 100, epoch)
+    writer.close()
+
+    test_probs = torch.cat([torch.stack(batch) for batch in class_probs])
+    test_preds = torch.cat(class_preds)
+    for i in range(len(classes)):
+        add_pr_curve_tensorboard(i, test_probs, test_preds, epoch)
+    #print(loss, total, correct, total)
+    return loss / total, correct / total
+
+
+def save_checkpoint(state, ckpt_dir='checkpoint'):
+    print("Checkpoint {} saved".format(state['epoch']))
+    filename = "ckpt_epoch" + str(state['epoch']) + ".pth"
+    filepath = os.path.join(ckpt_dir, filename)
+    torch.save(state, filepath)
+
+def adjust_lr(optimizer, epoch):
+    global state
+    new_lr = state["lr"] * (0.5 ** (epoch // 40)) if state["lr"] > 1e-7 else state["lr"]
+    if new_lr != state["lr"]:
+        state["lr"] = new_lr
+        print("Updating learning rate: {}".format(state["lr"]))
+        for param_group in optimizer.param_groups:
+            param_group["lr"] = state["lr"]
+
+def add_pr_curve_tensorboard(class_index, test_probs, test_preds, global_step=0):
+    global classes
+    '''
+    Takes in a "class_index" from 0 to 9 and plots the corresponding
+    precision-recall curve
+    '''
+    tensorboard_preds = test_preds == class_index
+    tensorboard_probs = test_probs[:, class_index]
+
+    writer.add_pr_curve(classes[class_index],
+                        tensorboard_preds,
+                        tensorboard_probs,
+                        global_step=global_step)
+    writer.close()
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file
diff --git a/cpp/ptq/training/vgg16/vgg16.py b/cpp/ptq/training/vgg16/vgg16.py
new file mode 100644
--- /dev/null
+++ b/cpp/ptq/training/vgg16/vgg16.py
@@ -0,0 +1,64 @@
+'''
+# Reference
+- [Very Deep Convolutional Networks for Large-Scale Image Recognition](
+    https://arxiv.org/abs/1409.1556) (ICLR 2015)
+'''
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from functools import reduce
+
+class VGG(nn.Module):
+    def __init__(self, layer_spec, num_classes=1000, init_weights=False):
+        super(VGG, self).__init__()
+
+        layers = []
+        in_channels = 3
+        for l in layer_spec:
+            if l == 'pool':
+                layers.append(nn.MaxPool2d(kernel_size=2, stride=2))
+            else:
+                layers += [
+                    nn.Conv2d(in_channels, l, kernel_size=3, padding=1),
+                    nn.BatchNorm2d(l),
+                    nn.ReLU()
+                ]
+                in_channels = l
+
+        self.features = nn.Sequential(*layers)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.classifier = nn.Sequential(
+            nn.Linear(512 * 1 * 1, 4096),
+            nn.ReLU(),
+            nn.Dropout(),
+            nn.Linear(4096, 4096),
+            nn.ReLU(),
+            nn.Dropout(),
+            nn.Linear(4096, num_classes)
+        )
+        if init_weights:
+            self._initialize_weights()
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.avgpool(x)
+        x = torch.flatten(x,1)
+        x = self.classifier(x)
+        return x
+
+def vgg16(num_classes=1000, init_weights=False):
+    vgg16_cfg = [64, 64, 'pool', 128, 128, 'pool', 256, 256, 256, 'pool', 512, 512, 512, 'pool', 512, 512, 512, 'pool']
+    return VGG(vgg16_cfg, num_classes, init_weights)
\ No newline at end of file


Test patch:
diff --git a/tests/BUILD b/tests/BUILD
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -5,4 +5,11 @@ test_suite(
         "//tests/modules:test_modules"
     ],
 )
-    
+
+test_suite(
+    name = "required_and_optional_tests",
+    tests = [
+        ":tests",
+        "//tests/accuracy:test_accuracy"
+    ]
+)
\ No newline at end of file
diff --git a/tests/accuracy/BUILD b/tests/accuracy/BUILD
new file mode 100644
--- /dev/null
+++ b/tests/accuracy/BUILD
@@ -0,0 +1,73 @@
+filegroup(
+    name = "jit_models",
+    srcs = glob(["**/*.jit.pt"])
+)
+
+test_suite(
+    name = "test_accuracy",
+    tests = [
+        ":test_int8_accuracy",
+        ":test_fp16_accuracy",
+        ":test_fp32_accuracy",
+    ]
+)
+
+cc_test(
+    name = "test_int8_accuracy",
+    srcs = ["test_int8_accuracy.cpp"],
+    deps = [
+        ":accuracy_test",
+        "//tests/accuracy/datasets:cifar10"
+    ],
+    data = [
+        ":jit_models",
+    ]
+)
+
+cc_test(
+    name = "test_fp16_accuracy",
+    srcs = ["test_fp16_accuracy.cpp"],
+    deps = [
+        ":accuracy_test",
+        "//tests/accuracy/datasets:cifar10"
+    ],
+    data = [
+        ":jit_models",
+    ]
+)
+
+cc_test(
+    name = "test_fp32_accuracy",
+    srcs = ["test_fp32_accuracy.cpp"],
+    deps = [
+        ":accuracy_test",
+        "//tests/accuracy/datasets:cifar10"
+    ],
+    data = [
+        ":jit_models",
+    ]
+)
+
+cc_binary(
+    name = "test",
+    srcs = ["test.cpp"],
+    deps = [
+        ":accuracy_test",
+        "//tests/accuracy/datasets:cifar10"
+    ],
+    data = [
+        ":jit_models",
+    ]
+)
+
+
+cc_library(
+    name = "accuracy_test",
+    hdrs = ["accuracy_test.h"],
+    deps = [
+        "//cpp/api:trtorch",
+        "//tests/util",
+        "@libtorch//:libtorch",
+        "@googletest//:gtest_main",
+    ],
+)
diff --git a/tests/accuracy/accuracy_test.h b/tests/accuracy/accuracy_test.h
new file mode 100644
--- /dev/null
+++ b/tests/accuracy/accuracy_test.h
@@ -0,0 +1,30 @@
+#include <utility>
+#include "torch/script.h"
+#include "gtest/gtest.h"
+#include "tests/util/util.h"
+#include "trtorch/trtorch.h"
+#include "c10/cuda/CUDACachingAllocator.h"
+
+// TODO: Extend this to support other datasets
+class AccuracyTests
+    : public testing::TestWithParam<std::string> {
+public:
+    void SetUp() override {
+        auto params = GetParam();
+        auto module_path = params;
+        try {
+            // Deserialize the ScriptModule from a file using torch::jit::load().
+            mod = torch::jit::load(module_path);
+        }
+        catch (const c10::Error& e) {
+            std::cerr << "error loading the model\n";
+            return;
+        }
+    }
+
+    void TearDown() {
+        c10::cuda::CUDACachingAllocator::emptyCache();
+    }
+protected:
+    torch::jit::script::Module mod;
+};
diff --git a/tests/accuracy/datasets/BUILD b/tests/accuracy/datasets/BUILD
new file mode 100644
--- /dev/null
+++ b/tests/accuracy/datasets/BUILD
@@ -0,0 +1,23 @@
+package(default_visibility = ["//visibility:public"])
+
+cc_library(
+    name = "cifar10",
+    hdrs = [
+        "cifar10.h"
+    ],
+    srcs = [
+        "cifar10.cpp"
+    ],
+    deps = [
+        "@libtorch//:libtorch"
+    ],
+    data = [
+        ":cifar10_data"
+    ]
+
+)
+
+filegroup(
+    name = "cifar10_data",
+    srcs = glob(["data/cifar-10-batches-bin/**/*.bin"])
+)
\ No newline at end of file
diff --git a/tests/accuracy/datasets/cifar10.cpp b/tests/accuracy/datasets/cifar10.cpp
new file mode 100644
--- /dev/null
+++ b/tests/accuracy/datasets/cifar10.cpp
@@ -0,0 +1,132 @@
+#include "tests/accuracy/datasets/cifar10.h"
+
+#include "torch/torch.h"
+#include "torch/data/example.h"
+#include "torch/types.h"
+
+#include <iostream>
+#include <cstddef>
+#include <fstream>
+#include <string>
+#include <vector>
+#include <utility>
+#include <sstream>
+#include <memory>
+
+namespace datasets {
+namespace {
+constexpr const char* kTrainFilenamePrefix = "data_batch_";
+constexpr const uint32_t kNumTrainFiles = 5;
+constexpr const char* kTestFilename = "test_batch.bin";
+constexpr const size_t kLabelSize = 1; // B
+constexpr const size_t kImageSize = 3072; // B
+constexpr const size_t kImageDim = 32;
+constexpr const size_t kImageChannels = 3;
+constexpr const size_t kBatchSize = 10000;
+
+std::pair<torch::Tensor, torch::Tensor> read_batch(const std::string& path) {
+    std::ifstream batch;
+    batch.open(path, std::ios::in|std::ios::binary|std::ios::ate);
+
+    auto file_size = batch.tellg();
+    std::unique_ptr<char[]> buf(new char[file_size]);
+
+    batch.seekg(0, std::ios::beg);
+    batch.read(buf.get(), file_size);
+    batch.close();
+
+    std::vector<uint8_t> labels;
+    std::vector<torch::Tensor> images;
+    labels.reserve(kBatchSize);
+    images.reserve(kBatchSize);
+
+    for (size_t i = 0; i < kBatchSize; i++) {
+        uint8_t label = buf[i * (kImageSize + kLabelSize)];
+        std::vector<uint8_t> image;
+        image.reserve(kImageSize);
+        std::copy(&buf[i * (kImageSize + kLabelSize) + 1], &buf[i * (kImageSize + kLabelSize) + kImageSize], std::back_inserter(image));
+        labels.push_back(label);
+        auto image_tensor = torch::from_blob(image.data(),
+                                             {kImageChannels, kImageDim, kImageDim},
+                                             torch::TensorOptions().dtype(torch::kU8)).to(torch::kF32);
+        images.push_back(image_tensor);
+    }
+
+    auto labels_tensor = torch::from_blob(labels.data(),
+                                          {kBatchSize},
+                                          torch::TensorOptions().dtype(torch::kU8)).to(torch::kF32);
+    assert(labels_tensor.size(0) == kBatchSize);
+
+    auto images_tensor = torch::stack(images);
+    assert(images_tensor.size(0) == kBatchSize);
+
+    return std::make_pair(images_tensor, labels_tensor);
+}
+
+std::pair<torch::Tensor, torch::Tensor> read_train_data(const std::string& root) {
+    std::vector<torch::Tensor> images, targets;
+    for(uint32_t i = 1; i <= 5; i++) {
+        std::stringstream ss;
+        ss << root << '/' << kTrainFilenamePrefix << i << ".bin";
+        auto batch = read_batch(ss.str());
+        images.push_back(batch.first);
+        targets.push_back(batch.second);
+    }
+
+    torch::Tensor image_tensor = std::accumulate(++images.begin(), images.end(), *images.begin(), [&](torch::Tensor a, torch::Tensor b) {return torch::cat({a, b}, 0);});
+    torch::Tensor target_tensor = std::accumulate(++targets.begin(), targets.end(), *targets.begin(), [&](torch::Tensor a, torch::Tensor b) {return torch::cat({a, b}, 0);});
+
+    return std::make_pair(image_tensor, target_tensor);
+}
+
+std::pair<torch::Tensor, torch::Tensor> read_test_data(const std::string& root) {
+    std::stringstream ss;
+    ss << root << '/' << kTestFilename;
+    return read_batch(ss.str());
+}
+}
+
+CIFAR10::CIFAR10(const std::string& root, Mode mode)
+    : mode_(mode) {
+
+    std::pair<torch::Tensor, torch::Tensor> data;
+    if (mode_ == Mode::kTrain) {
+        data = read_train_data(root);
+    } else {
+        data = read_test_data(root);
+    }
+
+    images_ = std::move(data.first);
+    targets_ = std::move(data.second);
+    assert(images_.sizes()[0] == images_.sizes()[0]);
+}
+
+torch::data::Example<> CIFAR10::get(size_t index) {
+    return {images_[index], targets_[index]};
+}
+
+c10::optional<size_t> CIFAR10::size() const {
+    return images_.size(0);
+}
+
+bool CIFAR10::is_train() const noexcept {
+    return mode_ == Mode::kTrain;
+}
+
+const torch::Tensor& CIFAR10::images() const {
+    return images_;
+}
+
+const torch::Tensor& CIFAR10::targets() const {
+    return targets_;
+}
+
+CIFAR10&& CIFAR10::use_subset(int64_t new_size) {
+    assert(new_size <= images_.sizes()[0]);
+    images_ = images_.slice(0, 0, new_size);
+    targets_ = targets_.slice(0, 0, new_size);
+    return std::move(*this);
+}
+
+} // namespace datasets
+
diff --git a/tests/accuracy/datasets/cifar10.h b/tests/accuracy/datasets/cifar10.h
new file mode 100644
--- /dev/null
+++ b/tests/accuracy/datasets/cifar10.h
@@ -0,0 +1,45 @@
+#pragma once
+
+#include "torch/data/datasets/base.h"
+#include "torch/data/example.h"
+#include "torch/types.h"
+
+#include <cstddef>
+#include <string>
+
+namespace datasets {
+// The CIFAR10 Dataset
+class CIFAR10 : public torch::data::datasets::Dataset<CIFAR10> {
+public:
+    // The mode in which the dataset is loaded
+    enum class Mode { kTrain, kTest };
+
+    // Loads CIFAR10 from un-tarred file
+    // Dataset can be found https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz
+    // Root path should be the directory that contains the content of tarball
+    explicit CIFAR10(const std::string& root, Mode mode = Mode::kTrain);
+
+    // Returns the pair at index in the dataset
+    torch::data::Example<> get(size_t index) override;
+
+    // The size of the dataset
+    c10::optional<size_t> size() const override;
+
+    // The mode the dataset is in
+    bool is_train() const noexcept;
+
+    // Returns all images stacked into a single tensor
+    const torch::Tensor& images() const;
+
+    // Returns all targets stacked into a single tensor
+    const torch::Tensor& targets() const;
+
+    // Trims the dataset to the first n pairs
+    CIFAR10&& use_subset(int64_t new_size);
+
+
+private:
+    Mode mode_;
+    torch::Tensor images_, targets_;
+};
+} // namespace datasets
diff --git a/tests/accuracy/test_fp16_accuracy.cpp b/tests/accuracy/test_fp16_accuracy.cpp
new file mode 100644
--- /dev/null
+++ b/tests/accuracy/test_fp16_accuracy.cpp
@@ -0,0 +1,56 @@
+#include "accuracy_test.h"
+#include "datasets/cifar10.h"
+#include "torch/torch.h"
+
+TEST_P(AccuracyTests, FP16AccuracyIsClose) {
+    auto eval_dataset = datasets::CIFAR10("tests/accuracy/datasets/data/cifar-10-batches-bin/", datasets::CIFAR10::Mode::kTest)
+                                    .use_subset(3200)
+                                    .map(torch::data::transforms::Normalize<>({0.4914, 0.4822, 0.4465},
+                                                                              {0.2023, 0.1994, 0.2010}))
+                                    .map(torch::data::transforms::Stack<>());
+    auto eval_dataloader = torch::data::make_data_loader(std::move(eval_dataset), torch::data::DataLoaderOptions()
+                                                                                                    .batch_size(32)
+                                                                                                    .workers(2));
+
+    // Check the FP32 accuracy in JIT
+    torch::Tensor jit_correct = torch::zeros({1}, {torch::kCUDA}), jit_total = torch::zeros({1}, {torch::kCUDA});
+    for (auto batch : *eval_dataloader) {
+        auto images = batch.data.to(torch::kCUDA);
+        auto targets = batch.target.to(torch::kCUDA);
+
+        auto outputs = mod.forward({images});
+        auto predictions = std::get<1>(torch::max(outputs.toTensor(), 1, false));
+
+        jit_total += targets.sizes()[0];
+        jit_correct += torch::sum(torch::eq(predictions, targets));
+    }
+    torch::Tensor jit_accuracy = jit_correct / jit_total;
+
+    std::vector<std::vector<int64_t>> input_shape = {{32, 3, 32, 32}};
+    auto extra_info = trtorch::ExtraInfo({input_shape});
+    extra_info.op_precision = torch::kF16;
+
+    auto trt_mod = trtorch::CompileGraph(mod, extra_info);
+
+    torch::Tensor trt_correct = torch::zeros({1}, {torch::kCUDA}), trt_total = torch::zeros({1}, {torch::kCUDA});
+    for (auto batch : *eval_dataloader) {
+        auto images = batch.data.to(torch::kCUDA).to(torch::kF16);
+        auto targets = batch.target.to(torch::kCUDA).to(torch::kF16);
+
+        auto outputs = trt_mod.forward({images});
+        auto predictions = std::get<1>(torch::max(outputs.toTensor(), 1, false));
+        predictions = predictions.reshape(predictions.sizes()[0]);
+
+        trt_total += targets.sizes()[0];
+        trt_correct += torch::sum(torch::eq(predictions, targets));
+    }
+
+    torch::Tensor trt_accuracy = trt_correct / trt_total;
+
+    ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_accuracy, trt_accuracy, 3));
+}
+
+
+INSTANTIATE_TEST_SUITE_P(FP16AccuracyIsCloseSuite,
+                         AccuracyTests,
+                         testing::Values("tests/accuracy/vgg16_cifar10.jit.pt"));
diff --git a/tests/accuracy/test_fp32_accuracy.cpp b/tests/accuracy/test_fp32_accuracy.cpp
new file mode 100644
--- /dev/null
+++ b/tests/accuracy/test_fp32_accuracy.cpp
@@ -0,0 +1,56 @@
+#include "accuracy_test.h"
+#include "datasets/cifar10.h"
+#include "torch/torch.h"
+
+TEST_P(AccuracyTests, FP16AccuracyIsClose) {
+    auto eval_dataset = datasets::CIFAR10("tests/accuracy/datasets/data/cifar-10-batches-bin/", datasets::CIFAR10::Mode::kTest)
+                                    .use_subset(3200)
+                                    .map(torch::data::transforms::Normalize<>({0.4914, 0.4822, 0.4465},
+                                                                              {0.2023, 0.1994, 0.2010}))
+                                    .map(torch::data::transforms::Stack<>());
+    auto eval_dataloader = torch::data::make_data_loader(std::move(eval_dataset), torch::data::DataLoaderOptions()
+                                                                                                    .batch_size(32)
+                                                                                                    .workers(2));
+
+    // Check the FP32 accuracy in JIT
+    torch::Tensor jit_correct = torch::zeros({1}, {torch::kCUDA}), jit_total = torch::zeros({1}, {torch::kCUDA});
+    for (auto batch : *eval_dataloader) {
+        auto images = batch.data.to(torch::kCUDA);
+        auto targets = batch.target.to(torch::kCUDA);
+
+        auto outputs = mod.forward({images});
+        auto predictions = std::get<1>(torch::max(outputs.toTensor(), 1, false));
+
+        jit_total += targets.sizes()[0];
+        jit_correct += torch::sum(torch::eq(predictions, targets));
+    }
+    torch::Tensor jit_accuracy = jit_correct / jit_total;
+
+    std::vector<std::vector<int64_t>> input_shape = {{32, 3, 32, 32}};
+    auto extra_info = trtorch::ExtraInfo({input_shape});
+    extra_info.op_precision = torch::kF32;
+
+    auto trt_mod = trtorch::CompileGraph(mod, extra_info);
+
+    torch::Tensor trt_correct = torch::zeros({1}, {torch::kCUDA}), trt_total = torch::zeros({1}, {torch::kCUDA});
+    for (auto batch : *eval_dataloader) {
+        auto images = batch.data.to(torch::kCUDA);
+        auto targets = batch.target.to(torch::kCUDA);
+
+        auto outputs = trt_mod.forward({images});
+        auto predictions = std::get<1>(torch::max(outputs.toTensor(), 1, false));
+        predictions = predictions.reshape(predictions.sizes()[0]);
+
+        trt_total += targets.sizes()[0];
+        trt_correct += torch::sum(torch::eq(predictions, targets));
+    }
+
+    torch::Tensor trt_accuracy = trt_correct / trt_total;
+
+    ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_accuracy, trt_accuracy, 3));
+}
+
+
+INSTANTIATE_TEST_SUITE_P(FP16AccuracyIsCloseSuite,
+                         AccuracyTests,
+                         testing::Values("tests/accuracy/vgg16_cifar10.jit.pt"));
diff --git a/tests/accuracy/test_int8_accuracy.cpp b/tests/accuracy/test_int8_accuracy.cpp
new file mode 100644
--- /dev/null
+++ b/tests/accuracy/test_int8_accuracy.cpp
@@ -0,0 +1,83 @@
+#include "accuracy_test.h"
+#include "datasets/cifar10.h"
+#include "torch/torch.h"
+
+TEST_P(AccuracyTests, FP16AccuracyIsClose) {
+    auto calibration_dataset = datasets::CIFAR10("tests/accuracy/datasets/data/cifar-10-batches-bin/", datasets::CIFAR10::Mode::kTest)
+                                    .use_subset(320)
+                                    .map(torch::data::transforms::Normalize<>({0.4914, 0.4822, 0.4465},
+                                                                              {0.2023, 0.1994, 0.2010}))
+                                    .map(torch::data::transforms::Stack<>());
+    auto calibration_dataloader = torch::data::make_data_loader(std::move(calibration_dataset), torch::data::DataLoaderOptions()
+                                                                                                    .batch_size(32)
+                                                                                                    .workers(2));
+
+    std::string calibration_cache_file = "/tmp/vgg16_TRT_ptq_calibration.cache";
+
+    auto calibrator = trtorch::ptq::make_int8_calibrator(std::move(calibration_dataloader), calibration_cache_file, true);
+    //auto calibrator = trtorch::ptq::make_int8_cache_calibrator(calibration_cache_file);
+
+
+    std::vector<std::vector<int64_t>> input_shape = {{32, 3, 32, 32}};
+    // Configure settings for compilation
+    auto extra_info = trtorch::ExtraInfo({input_shape});
+    // Set operating precision to INT8
+    extra_info.op_precision = torch::kI8;
+    // Use the TensorRT Entropy Calibrator
+    extra_info.ptq_calibrator = calibrator;
+    // Set max batch size for the engine
+    extra_info.max_batch_size = 32;
+    // Set a larger workspace
+    extra_info.workspace_size = 1 << 28;
+
+    mod.eval();
+
+    // Dataloader moved into calibrator so need another for inference
+    auto eval_dataset = datasets::CIFAR10("tests/accuracy/datasets/data/cifar-10-batches-bin/", datasets::CIFAR10::Mode::kTest)
+                                    .use_subset(3200)
+                                    .map(torch::data::transforms::Normalize<>({0.4914, 0.4822, 0.4465},
+                                                                              {0.2023, 0.1994, 0.2010}))
+                                    .map(torch::data::transforms::Stack<>());
+    auto eval_dataloader = torch::data::make_data_loader(std::move(eval_dataset), torch::data::DataLoaderOptions()
+                                                                                                    .batch_size(32)
+                                                                                                    .workers(2));
+
+    // Check the FP32 accuracy in JIT
+    torch::Tensor jit_correct = torch::zeros({1}, {torch::kCUDA}), jit_total = torch::zeros({1}, {torch::kCUDA});
+    for (auto batch : *eval_dataloader) {
+        auto images = batch.data.to(torch::kCUDA);
+        auto targets = batch.target.to(torch::kCUDA);
+
+        auto outputs = mod.forward({images});
+        auto predictions = std::get<1>(torch::max(outputs.toTensor(), 1, false));
+
+        jit_total += targets.sizes()[0];
+        jit_correct += torch::sum(torch::eq(predictions, targets));
+    }
+    torch::Tensor jit_accuracy = jit_correct / jit_total;
+
+    // Compile Graph
+    auto trt_mod = trtorch::CompileGraph(mod, extra_info);
+
+    // Check the INT8 accuracy in TRT
+    torch::Tensor trt_correct = torch::zeros({1}, {torch::kCUDA}), trt_total = torch::zeros({1}, {torch::kCUDA});
+    for (auto batch : *eval_dataloader) {
+        auto images = batch.data.to(torch::kCUDA);
+        auto targets = batch.target.to(torch::kCUDA);
+
+        auto outputs = trt_mod.forward({images});
+        auto predictions = std::get<1>(torch::max(outputs.toTensor(), 1, false));
+        predictions = predictions.reshape(predictions.sizes()[0]);
+
+        trt_total += targets.sizes()[0];
+        trt_correct += torch::sum(torch::eq(predictions, targets)).item().toFloat();
+    }
+    torch::Tensor trt_accuracy = trt_correct / trt_total;
+
+    ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_accuracy, trt_accuracy, 3));
+}
+
+
+INSTANTIATE_TEST_SUITE_P(FP16AccuracyIsCloseSuite,
+                         AccuracyTests,
+                         testing::Values("tests/accuracy/vgg16_cifar10.jit.pt"));
\ No newline at end of file
diff --git a/tests/modules/BUILD b/tests/modules/BUILD
--- a/tests/modules/BUILD
+++ b/tests/modules/BUILD
@@ -23,7 +23,6 @@ cc_test(
     ]
 )
 
-
 cc_test(
     name = "test_modules_as_engines",
     srcs = ["test_modules_as_engines.cpp"],
diff --git a/tests/util/run_forward.cpp b/tests/util/run_forward.cpp
--- a/tests/util/run_forward.cpp
+++ b/tests/util/run_forward.cpp
@@ -18,7 +18,7 @@ std::vector<at::Tensor> RunModuleForwardAsEngine(torch::jit::script::Module& mod
     for (auto in : inputs) {
         input_ranges.push_back(in.sizes());
     }
-    
+
     auto engine = trtorch::ConvertGraphToTRTEngine(mod, "forward", input_ranges);
     return RunEngine(engine, inputs);
 }