From c8e26afd67da1a41658cbcb2d9f577b026620d31 Mon Sep 17 00:00:00 2001
From: Colin <colin@deve.work>
Date: Mon, 10 Jan 2022 08:05:23 +0000
Subject: [PATCH] Add AMP

---
 Amp/test.py    | 153 +++++++++++++
 Amp/trace.json | 607 +++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 760 insertions(+)
 create mode 100644 Amp/test.py
 create mode 100644 Amp/trace.json

diff --git a/Amp/test.py b/Amp/test.py
new file mode 100644
index 0000000..d9f58f8
--- /dev/null
+++ b/Amp/test.py
@@ -0,0 +1,153 @@
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import torchvision
+from torchvision import datasets, transforms
+import torchvision.models as models
+from torch.cuda.amp import autocast as autocast
+from torch.utils.data import Dataset, DataLoader
+import numpy as np
+import random
+import cv2
+from torch.profiler import profile, record_function, ProfilerActivity
+
+
+
+batchsize = 12
+
+
+class ManualDataset(Dataset):
+    def __init__(self, transform=None):
+        self.transform = transform
+
+    def __getitem__(self, index):
+
+        # data = np.random.random_integers(0,255,(28,28))
+        data = np.zeros((28, 28), dtype="float32")
+        data = data.astype("float32")
+
+        radiu = int(random.random()*7+3)
+        basex = random.randint(radiu, 29-radiu)
+        basey = random.randint(radiu, 29-radiu)
+        ifcircle = random.random()
+
+        if ifcircle > 0.5:
+            cv2.circle(data, (basex, basey), radiu, 255, -1)
+            angle = random.random() * 360
+            M = cv2.getRotationMatrix2D((basex, basey), angle, 1.0)
+            data = cv2.warpAffine(data, M, (28, 28))
+            label = 0
+        else:
+            cv2.rectangle(data, (basex-radiu, basey-radiu),
+                          (basex+radiu, basey+radiu), 255, -1)
+            angle = random.random() * 360
+            M = cv2.getRotationMatrix2D((basex, basey), angle, 1.0)
+            data = cv2.warpAffine(data, M, (28, 28))
+            label = 1
+
+        # cv2.imwrite("test.jpg",data)
+
+        data = (data - 128) / 256.0
+        img = torch.from_numpy(data)
+        img = img.view(1, 28, 28)
+        return img, label
+
+    def __len__(self):
+        return 10000
+
+
+train_loader = torch.utils.data.DataLoader(
+    ManualDataset(transform=transforms.Compose([
+        # transforms.ColorJitter(0.2,0.2),
+        # transforms.RandomRotation(30),
+        # transforms.RandomResizedCrop(28),
+        transforms.ToTensor(),
+        transforms.Normalize((128,), (256,)), ])),
+    batch_size=batchsize, shuffle=True, num_workers=2)
+test_loader = torch.utils.data.DataLoader(
+    ManualDataset(transform=transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize((128,), (256,))])),
+    batch_size=batchsize, shuffle=True, num_workers=2)
+
+
+class Net(nn.Module):
+    def __init__(self):
+        super(Net, self).__init__()
+        layers = []
+        layers += [nn.Conv2d(1,  8, kernel_size=5),
+                   nn.MaxPool2d(kernel_size=2, stride=2), nn.Sigmoid()]
+        layers += [nn.Conv2d(8,  8, kernel_size=3),
+                   nn.MaxPool2d(kernel_size=2, stride=2), nn.Sigmoid()]
+        layers += [nn.Conv2d(8,  10, kernel_size=5)]
+        self.features = nn.Sequential(*layers)
+
+        # self.conv1 = nn.Conv2d(1,  8, kernel_size=5)
+        # self.conv2 = self.__conv(8,  8, kernel_size=3)
+        # self.conv3 = self.__conv(8, 10, kernel_size=5)
+
+    def forward(self, x):
+        # x = F.sigmoid(F.max_pool2d(self.conv1(x), 2))
+        # x = F.sigmoid(F.max_pool2d(self.conv2(x), 2))
+        # x = self.conv3(x)
+        x = self.features(x)
+        x = x.view(-1, 1*10)
+
+        return F.log_softmax(x, dim=1)
+
+
+model = Net()
+# model = model.cuda()
+
+optimizer = optim.SGD(model.parameters(), lr=0.0005)
+
+for batch_idx, (data, target) in enumerate(train_loader):
+    # data = data.cuda()
+    # target = target.cuda()
+    optimizer.zero_grad()
+
+    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], with_stack=True, record_shapes=True) as prof:
+        with record_function("model_inference"):
+            with autocast():
+                output = model(data)
+                loss = F.nll_loss(output, target)
+    # print(prof.key_averages().table(sort_by="cpu_time_total"))
+    print(prof.key_averages().table(sort_by="cpu_time_total"))
+    prof.export_chrome_trace("trace.json")
+
+    exit()
+    # 反向传播在autocast上下文之外
+    loss.backward()
+    optimizer.step()
+
+
+# from torch.cuda.amp import autocast as autocast
+
+# # 创建model，默认是torch.FloatTensor
+# model = Net().cuda()
+# optimizer = optim.SGD(model.parameters(), ...)
+
+# # 在训练最开始之前实例化一个GradScaler对象
+# scaler = GradScaler()
+
+# for epoch in epochs:
+#     for input, target in data:
+#         optimizer.zero_grad()
+
+#         # 前向过程(model + loss)开启 autocast
+#         with autocast():
+#             output = model(input)
+#             loss = loss_fn(output, target)
+
+#         # Scales loss. 为了梯度放大.
+#         scaler.scale(loss).backward()
+
+#         # scaler.step() 首先把梯度的值unscale回来.
+#         # 如果梯度的值不是 infs 或者 NaNs, 那么调用optimizer.step()来更新权重,
+#         # 否则，忽略step调用，从而保证权重不更新（不被破坏）
+#         scaler.step(optimizer)
+
+#         # 准备着，查看是否要增大scaler
+#         scaler.update()
diff --git a/Amp/trace.json b/Amp/trace.json
new file mode 100644
index 0000000..6b53e34
--- /dev/null
+++ b/Amp/trace.json
@@ -0,0 +1,607 @@
+
+{
+  "schemaVersion": 1,
+  "deviceProperties": [
+    {
+      "id": 0, "name": "NVIDIA GeForce RTX 2070 SUPER", "totalGlobalMem": 8367439872,
+      "computeMajor": 7, "computeMinor": 5,
+      "maxThreadsPerBlock": 1024, "maxThreadsPerMultiprocessor": 1024,
+      "regsPerBlock": 65536, "regsPerMultiprocessor": 65536, "warpSize": 32,
+      "sharedMemPerBlock": 49152, "sharedMemPerMultiprocessor": 65536,
+      "numSms": 40, "sharedMemPerBlockOptin": 65536
+    }
+  ],
+  "traceEvents": [
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::empty", "pid": 14436, "tid": 14436,
+    "ts": 1641800282593594, "dur": 7,
+    "args": {
+       "External id": 2,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[], [], [], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/autograd/profiler.py(432): __init__;test.py(112): <module>;", "Input type": ["GenericList", "Int", "", "", "Bool", ""]
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::zero_", "pid": 14436, "tid": 14436,
+    "ts": 1641800282593614, "dur": 1,
+    "args": {
+       "External id": 3,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[1]], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/autograd/profiler.py(432): __init__;test.py(112): <module>;", "Input type": ["float"]
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::zeros", "pid": 14436, "tid": 14436,
+    "ts": 1641800282593579, "dur": 38,
+    "args": {
+       "External id": 1,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[], [], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/autograd/profiler.py(432): __init__;test.py(112): <module>;", "Input type": ["GenericList", "Int", "", "", "Bool"]
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::empty", "pid": 14436, "tid": 14436,
+    "ts": 1641800282593648, "dur": 1,
+    "args": {
+       "External id": 5,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[], [], [], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/autograd/profiler.py(435): __enter__;test.py(112): <module>;", "Input type": ["GenericList", "Int", "", "", "", ""]
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::empty", "pid": 14436, "tid": 14436,
+    "ts": 1641800282602596, "dur": 6,
+    "args": {
+       "External id": 10,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[], [], [], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(442): _conv_forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(446): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["GenericList", "Int", "Int", "", "", ""]
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::as_strided_", "pid": 14436, "tid": 14436,
+    "ts": 1641800282602862, "dur": 4,
+    "args": {
+       "External id": 11,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 8, 24, 24], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(442): _conv_forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(446): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float", "GenericList", "GenericList", ""]
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::mkldnn_convolution", "pid": 14436, "tid": 14436,
+    "ts": 1641800282593898, "dur": 8986,
+    "args": {
+       "External id": 9,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 1, 28, 28], [8, 1, 5, 5], [8], [], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(442): _conv_forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(446): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float", "float", "float", "GenericList", "GenericList", "GenericList", "Int"], "Fwd thread id": 0, "Sequence number": 0
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::_convolution", "pid": 14436, "tid": 14436,
+    "ts": 1641800282593879, "dur": 9007,
+    "args": {
+       "External id": 8,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 1, 28, 28], [8, 1, 5, 5], [8], [], [], [], [], [], [], [], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(442): _conv_forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(446): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float", "float", "float", "GenericList", "GenericList", "GenericList", "Bool", "GenericList", "Int", "Bool", "Bool", "Bool", "Bool"], "Fwd thread id": 0, "Sequence number": 0
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::convolution", "pid": 14436, "tid": 14436,
+    "ts": 1641800282593860, "dur": 9027,
+    "args": {
+       "External id": 7,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 1, 28, 28], [8, 1, 5, 5], [8], [], [], [], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(442): _conv_forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(446): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float", "float", "float", "GenericList", "GenericList", "GenericList", "Bool", "GenericList", "Int"], "Fwd thread id": 0, "Sequence number": 0
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::conv2d", "pid": 14436, "tid": 14436,
+    "ts": 1641800282593839, "dur": 9051,
+    "args": {
+       "External id": 6,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 1, 28, 28], [8, 1, 5, 5], [8], [], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(442): _conv_forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(446): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float", "float", "float", "GenericList", "GenericList", "GenericList", "Int"], "Fwd thread id": 0, "Sequence number": 0
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::max_pool2d_with_indices", "pid": 14436, "tid": 14436,
+    "ts": 1641800282603040, "dur": 428,
+    "args": {
+       "External id": 13,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 8, 24, 24], [], [], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/functional.py(719): _max_pool2d;/home/colin/anaconda3/lib/python3.8/site-packages/torch/_jit_internal.py(422): fn;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/pooling.py(162): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float", "GenericList", "GenericList", "GenericList", "GenericList", "Bool"], "Fwd thread id": 0, "Sequence number": 1
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::max_pool2d", "pid": 14436, "tid": 14436,
+    "ts": 1641800282603015, "dur": 461,
+    "args": {
+       "External id": 12,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 8, 24, 24], [], [], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/functional.py(719): _max_pool2d;/home/colin/anaconda3/lib/python3.8/site-packages/torch/_jit_internal.py(422): fn;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/pooling.py(162): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float", "GenericList", "GenericList", "GenericList", "GenericList", "Bool"], "Fwd thread id": 0, "Sequence number": 1
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::sigmoid", "pid": 14436, "tid": 14436,
+    "ts": 1641800282603523, "dur": 365,
+    "args": {
+       "External id": 14,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 8, 12, 12]], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/activation.py(291): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float"], "Fwd thread id": 0, "Sequence number": 2
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::empty", "pid": 14436, "tid": 14436,
+    "ts": 1641800282604168, "dur": 2,
+    "args": {
+       "External id": 19,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[], [], [], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(442): _conv_forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(446): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["GenericList", "Int", "Int", "", "", ""]
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::as_strided_", "pid": 14436, "tid": 14436,
+    "ts": 1641800282604218, "dur": 0,
+    "args": {
+       "External id": 20,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 8, 10, 10], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(442): _conv_forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(446): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float", "GenericList", "GenericList", ""]
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::mkldnn_convolution", "pid": 14436, "tid": 14436,
+    "ts": 1641800282603942, "dur": 280,
+    "args": {
+       "External id": 18,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 8, 12, 12], [8, 8, 3, 3], [8], [], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(442): _conv_forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(446): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float", "float", "float", "GenericList", "GenericList", "GenericList", "Int"], "Fwd thread id": 0, "Sequence number": 3
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::_convolution", "pid": 14436, "tid": 14436,
+    "ts": 1641800282603931, "dur": 293,
+    "args": {
+       "External id": 17,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 8, 12, 12], [8, 8, 3, 3], [8], [], [], [], [], [], [], [], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(442): _conv_forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(446): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float", "float", "float", "GenericList", "GenericList", "GenericList", "Bool", "GenericList", "Int", "Bool", "Bool", "Bool", "Bool"], "Fwd thread id": 0, "Sequence number": 3
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::convolution", "pid": 14436, "tid": 14436,
+    "ts": 1641800282603921, "dur": 308,
+    "args": {
+       "External id": 16,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 8, 12, 12], [8, 8, 3, 3], [8], [], [], [], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(442): _conv_forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(446): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float", "float", "float", "GenericList", "GenericList", "GenericList", "Bool", "GenericList", "Int"], "Fwd thread id": 0, "Sequence number": 3
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::conv2d", "pid": 14436, "tid": 14436,
+    "ts": 1641800282603911, "dur": 319,
+    "args": {
+       "External id": 15,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 8, 12, 12], [8, 8, 3, 3], [8], [], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(442): _conv_forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(446): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float", "float", "float", "GenericList", "GenericList", "GenericList", "Int"], "Fwd thread id": 0, "Sequence number": 3
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::max_pool2d_with_indices", "pid": 14436, "tid": 14436,
+    "ts": 1641800282604261, "dur": 19,
+    "args": {
+       "External id": 22,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 8, 10, 10], [], [], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/functional.py(719): _max_pool2d;/home/colin/anaconda3/lib/python3.8/site-packages/torch/_jit_internal.py(422): fn;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/pooling.py(162): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float", "GenericList", "GenericList", "GenericList", "GenericList", "Bool"], "Fwd thread id": 0, "Sequence number": 4
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::max_pool2d", "pid": 14436, "tid": 14436,
+    "ts": 1641800282604251, "dur": 31,
+    "args": {
+       "External id": 21,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 8, 10, 10], [], [], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/functional.py(719): _max_pool2d;/home/colin/anaconda3/lib/python3.8/site-packages/torch/_jit_internal.py(422): fn;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/pooling.py(162): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float", "GenericList", "GenericList", "GenericList", "GenericList", "Bool"], "Fwd thread id": 0, "Sequence number": 4
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::sigmoid", "pid": 14436, "tid": 14436,
+    "ts": 1641800282604295, "dur": 8,
+    "args": {
+       "External id": 23,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 8, 5, 5]], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/activation.py(291): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float"], "Fwd thread id": 0, "Sequence number": 5
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::empty", "pid": 14436, "tid": 14436,
+    "ts": 1641800282605085, "dur": 2,
+    "args": {
+       "External id": 28,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[], [], [], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(442): _conv_forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(446): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["GenericList", "Int", "Int", "", "", ""]
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::as_strided_", "pid": 14436, "tid": 14436,
+    "ts": 1641800282606358, "dur": 1,
+    "args": {
+       "External id": 29,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 10, 1, 1], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(442): _conv_forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(446): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float", "GenericList", "GenericList", ""]
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::mkldnn_convolution", "pid": 14436, "tid": 14436,
+    "ts": 1641800282604356, "dur": 2008,
+    "args": {
+       "External id": 27,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 8, 5, 5], [10, 8, 5, 5], [10], [], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(442): _conv_forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(446): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float", "float", "float", "GenericList", "GenericList", "GenericList", "Int"], "Fwd thread id": 0, "Sequence number": 6
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::_convolution", "pid": 14436, "tid": 14436,
+    "ts": 1641800282604344, "dur": 2021,
+    "args": {
+       "External id": 26,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 8, 5, 5], [10, 8, 5, 5], [10], [], [], [], [], [], [], [], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(442): _conv_forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(446): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float", "float", "float", "GenericList", "GenericList", "GenericList", "Bool", "GenericList", "Int", "Bool", "Bool", "Bool", "Bool"], "Fwd thread id": 0, "Sequence number": 6
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::convolution", "pid": 14436, "tid": 14436,
+    "ts": 1641800282604333, "dur": 2034,
+    "args": {
+       "External id": 25,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 8, 5, 5], [10, 8, 5, 5], [10], [], [], [], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(442): _conv_forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(446): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float", "float", "float", "GenericList", "GenericList", "GenericList", "Bool", "GenericList", "Int"], "Fwd thread id": 0, "Sequence number": 6
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::conv2d", "pid": 14436, "tid": 14436,
+    "ts": 1641800282604322, "dur": 2046,
+    "args": {
+       "External id": 24,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 8, 5, 5], [10, 8, 5, 5], [10], [], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(442): _conv_forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py(446): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py(141): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(95): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float", "float", "float", "GenericList", "GenericList", "GenericList", "Int"], "Fwd thread id": 0, "Sequence number": 6
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::view", "pid": 14436, "tid": 14436,
+    "ts": 1641800282606400, "dur": 13,
+    "args": {
+       "External id": 30,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 10, 1, 1], []], "Call stack": "test.py(96): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float", "GenericList"], "Fwd thread id": 0, "Sequence number": 7
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::_log_softmax", "pid": 14436, "tid": 14436,
+    "ts": 1641800282606479, "dur": 147,
+    "args": {
+       "External id": 32,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 10], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/functional.py(1769): log_softmax;test.py(98): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float", "Int", "Bool"], "Fwd thread id": 0, "Sequence number": 8
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::log_softmax", "pid": 14436, "tid": 14436,
+    "ts": 1641800282606463, "dur": 164,
+    "args": {
+       "External id": 31,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 10], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/functional.py(1769): log_softmax;test.py(98): forward;/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl;test.py(114): <module>;", "Input type": ["float", "Int", ""], "Fwd thread id": 0, "Sequence number": 8
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::nll_loss_forward", "pid": 14436, "tid": 14436,
+    "ts": 1641800282606769, "dur": 15,
+    "args": {
+       "External id": 35,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 10], [12], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/functional.py(2532): nll_loss;test.py(115): <module>;", "Input type": ["float", "long int", "", "Int", "Int"], "Fwd thread id": 0, "Sequence number": 9
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::nll_loss", "pid": 14436, "tid": 14436,
+    "ts": 1641800282606751, "dur": 34,
+    "args": {
+       "External id": 34,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 10], [12], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/functional.py(2532): nll_loss;test.py(115): <module>;", "Input type": ["float", "long int", "", "Int", "Int"], "Fwd thread id": 0, "Sequence number": 9
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "aten::nll_loss_nd", "pid": 14436, "tid": 14436,
+    "ts": 1641800282606731, "dur": 63,
+    "args": {
+       "External id": 33,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Input Dims": [[12, 10], [12], [], [], []], "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/nn/functional.py(2532): nll_loss;test.py(115): <module>;", "Input type": ["float", "long int", "", "Int", "Int"], "Fwd thread id": 0, "Sequence number": 9
+    }
+  },
+  {
+    "ph": "X", "cat": "cpu_op", 
+    "name": "model_inference", "pid": 14436, "tid": 14436,
+    "ts": 1641800282593642, "dur": 13169,
+    "args": {
+       "External id": 4,
+       "Trace name": "PyTorch Profiler", "Trace iteration": 0,
+       "Call stack": "/home/colin/anaconda3/lib/python3.8/site-packages/torch/autograd/profiler.py(435): __enter__;test.py(112): <module>;"
+    }
+  },
+  {
+    "ph": "X", "cat": "Runtime", 
+    "name": "cudaGetDeviceCount", "pid": 14436, "tid": 14436,
+    "ts": 1641800282606912, "dur": 328,
+    "args": {
+      "cbid": 3, "correlation": 1,
+      "external id": 0, "external ts": 0
+    }
+  },
+  {
+    "ph": "X", "cat": "Runtime", 
+    "name": "cudaGetDeviceProperties", "pid": 14436, "tid": 14436,
+    "ts": 1641800282607278, "dur": 691,
+    "args": {
+      "cbid": 4, "correlation": 4,
+      "external id": 0, "external ts": 0
+    }
+  },
+  {
+    "ph": "X", "cat": "Runtime", 
+    "name": "cudaGetDeviceCount", "pid": 14436, "tid": 14436,
+    "ts": 1641800282607978, "dur": 0,
+    "args": {
+      "cbid": 3, "correlation": 6,
+      "external id": 0, "external ts": 0
+    }
+  },
+  {
+    "ph": "X", "cat": "Runtime", 
+    "name": "cudaGetDeviceProperties", "pid": 14436, "tid": 14436,
+    "ts": 1641800282607980, "dur": 117,
+    "args": {
+      "cbid": 4, "correlation": 7,
+      "external id": 0, "external ts": 0
+    }
+  },
+  {
+    "ph": "X", "cat": "Runtime", 
+    "name": "cudaDeviceSynchronize", "pid": 14436, "tid": 14436,
+    "ts": 1641800282608432, "dur": 3091248,
+    "args": {
+      "cbid": 165, "correlation": 10,
+      "external id": 0, "external ts": 0
+    }
+  },
+  {
+    "name": "process_name", "ph": "M", "ts": 1641800282593491, "pid": 14436, "tid": 0,
+    "args": {
+      "name": "python"
+    }
+  },
+  {
+    "name": "process_labels", "ph": "M", "ts": 1641800282593491, "pid": 14436, "tid": 0,
+    "args": {
+      "labels": "CPU"
+    }
+  },
+  {
+    "name": "process_sort_index", "ph": "M", "ts": 1641800282593491, "pid": 14436, "tid": 0,
+    "args": {
+      "sort_index": 14436
+    }
+  },
+  {
+    "name": "process_name", "ph": "M", "ts": 1641800282593491, "pid": 0, "tid": 0,
+    "args": {
+      "name": "python"
+    }
+  },
+  {
+    "name": "process_labels", "ph": "M", "ts": 1641800282593491, "pid": 0, "tid": 0,
+    "args": {
+      "labels": "GPU 0"
+    }
+  },
+  {
+    "name": "process_sort_index", "ph": "M", "ts": 1641800282593491, "pid": 0, "tid": 0,
+    "args": {
+      "sort_index": 16777216
+    }
+  },
+  {
+    "name": "process_name", "ph": "M", "ts": 1641800282593491, "pid": 1, "tid": 0,
+    "args": {
+      "name": "python"
+    }
+  },
+  {
+    "name": "process_labels", "ph": "M", "ts": 1641800282593491, "pid": 1, "tid": 0,
+    "args": {
+      "labels": "GPU 1"
+    }
+  },
+  {
+    "name": "process_sort_index", "ph": "M", "ts": 1641800282593491, "pid": 1, "tid": 0,
+    "args": {
+      "sort_index": 16777217
+    }
+  },
+  {
+    "name": "process_name", "ph": "M", "ts": 1641800282593491, "pid": 2, "tid": 0,
+    "args": {
+      "name": "python"
+    }
+  },
+  {
+    "name": "process_labels", "ph": "M", "ts": 1641800282593491, "pid": 2, "tid": 0,
+    "args": {
+      "labels": "GPU 2"
+    }
+  },
+  {
+    "name": "process_sort_index", "ph": "M", "ts": 1641800282593491, "pid": 2, "tid": 0,
+    "args": {
+      "sort_index": 16777218
+    }
+  },
+  {
+    "name": "process_name", "ph": "M", "ts": 1641800282593491, "pid": 3, "tid": 0,
+    "args": {
+      "name": "python"
+    }
+  },
+  {
+    "name": "process_labels", "ph": "M", "ts": 1641800282593491, "pid": 3, "tid": 0,
+    "args": {
+      "labels": "GPU 3"
+    }
+  },
+  {
+    "name": "process_sort_index", "ph": "M", "ts": 1641800282593491, "pid": 3, "tid": 0,
+    "args": {
+      "sort_index": 16777219
+    }
+  },
+  {
+    "name": "process_name", "ph": "M", "ts": 1641800282593491, "pid": 4, "tid": 0,
+    "args": {
+      "name": "python"
+    }
+  },
+  {
+    "name": "process_labels", "ph": "M", "ts": 1641800282593491, "pid": 4, "tid": 0,
+    "args": {
+      "labels": "GPU 4"
+    }
+  },
+  {
+    "name": "process_sort_index", "ph": "M", "ts": 1641800282593491, "pid": 4, "tid": 0,
+    "args": {
+      "sort_index": 16777220
+    }
+  },
+  {
+    "name": "process_name", "ph": "M", "ts": 1641800282593491, "pid": 5, "tid": 0,
+    "args": {
+      "name": "python"
+    }
+  },
+  {
+    "name": "process_labels", "ph": "M", "ts": 1641800282593491, "pid": 5, "tid": 0,
+    "args": {
+      "labels": "GPU 5"
+    }
+  },
+  {
+    "name": "process_sort_index", "ph": "M", "ts": 1641800282593491, "pid": 5, "tid": 0,
+    "args": {
+      "sort_index": 16777221
+    }
+  },
+  {
+    "name": "process_name", "ph": "M", "ts": 1641800282593491, "pid": 6, "tid": 0,
+    "args": {
+      "name": "python"
+    }
+  },
+  {
+    "name": "process_labels", "ph": "M", "ts": 1641800282593491, "pid": 6, "tid": 0,
+    "args": {
+      "labels": "GPU 6"
+    }
+  },
+  {
+    "name": "process_sort_index", "ph": "M", "ts": 1641800282593491, "pid": 6, "tid": 0,
+    "args": {
+      "sort_index": 16777222
+    }
+  },
+  {
+    "name": "process_name", "ph": "M", "ts": 1641800282593491, "pid": 7, "tid": 0,
+    "args": {
+      "name": "python"
+    }
+  },
+  {
+    "name": "process_labels", "ph": "M", "ts": 1641800282593491, "pid": 7, "tid": 0,
+    "args": {
+      "labels": "GPU 7"
+    }
+  },
+  {
+    "name": "process_sort_index", "ph": "M", "ts": 1641800282593491, "pid": 7, "tid": 0,
+    "args": {
+      "sort_index": 16777223
+    }
+  },
+  {
+    "name": "thread_name", "ph": "M", "ts": 1641800282593491, "pid": 14436, "tid": 14436,
+    "args": {
+      "name": "thread 14436 (python)"
+    }
+  },
+  {
+    "name": "thread_sort_index", "ph": "M", "ts": 1641800282593491, "pid": 14436, "tid": 14436,
+    "args": {
+      "sort_index": 14436
+    }
+  },
+  {
+    "ph": "X", "cat": "Trace", "ts": 1641800282593491, "dur": 3106207,
+    "pid": "Spans", "tid": "PyTorch Profiler",
+    "name": "PyTorch Profiler (0)",
+    "args": {
+      "Op count": 0
+    }
+  },
+  {
+    "name": "process_sort_index", "ph": "M", "ts": 1641800282593491,
+    "pid": "Spans", "tid": 0,
+    "args": {
+      "sort_index": 536870912
+    }
+  },
+  {
+    "name": "Iteration Start: PyTorch Profiler", "ph": "i", "s": "g",
+    "pid": "Traces", "tid": "Trace PyTorch Profiler", "ts": 1641800282593491
+  },
+  {
+    "name": "Record Window End", "ph": "i", "s": "g",
+    "pid": "", "tid": "", "ts": 1641800285699878
+  }
+]}
\ No newline at end of file