FP16

Author: harrydobbs

debug/animation.py

@@ -0,0 +1,38 @@
+from manim import *
+
+
+class CreateCircle(Scene):
+    def construct(self):
+        self.camera.background_color = WHITE
+
+        tree_polygon_xyz = [
+            [4, 1, 0],  # middle right
+            [4, -2.5, 0],  # bottom right
+            [0, -2.5, 0],  # bottom left
+            [0, 3, 0],  # top left
+            [2, 1, 0],  # middle
+            [4, 3, 0],  # top right
+        ]
+
+        colorList = [RED, GREEN, BLUE, YELLOW]
+        for i in range(200):
+            point = Point(
+                location=[
+                    0.63 * np.random.randint(-4, 4),
+                    0.37 * np.random.randint(-4, 4),
+                    0,
+                ],
+                color=np.random.choice(colorList),
+            )
+            self.add(point)
+        for i in range(200):
+            point = Point(
+                location=[
+                    0.37 * np.random.randint(-4, 4),
+                    0.63 * np.random.randint(-4, 4),
+                    0,
+                ],
+                color=np.random.choice(colorList),
+            )
+            self.add(point)
+        self.add(point)

debug/fp16-2.py

@@ -0,0 +1,194 @@
+# Copyright 2021 Yan Yan
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import print_function
+import argparse
+import torch
+import spconv.pytorch as spconv
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.optim.lr_scheduler import StepLR
+import contextlib
+import torch.cuda.amp
+import time
+
+
+@contextlib.contextmanager
+def identity_ctx():
+    yield
+
+
+class Net(nn.Module):
+    def __init__(self):
+        super(Net, self).__init__()
+        self.net = spconv.SparseSequential(
+            # nn.BatchNorm1d(1),
+            spconv.SubMConv2d(1, 32, 3, 1),
+        )
+        self.fc1 = nn.Linear(14 * 14 * 64, 128)
+        self.fc2 = nn.Linear(128, 10)
+        self.dropout1 = nn.Dropout2d(0.25)
+        self.dropout2 = nn.Dropout2d(0.5)
+
+    def forward(self, x: torch.Tensor):
+        # x: [N, 28, 28, 1], must be NHWC tensor
+        x_sp = spconv.SparseConvTensor.from_dense(x.reshape(-1, 28, 28, 1))
+        x = self.net(x_sp)
+        return torch.flatten(x.features, 1)
+
+
+def train(args, model, device, train_loader, optimizer, epoch):
+    model.train()
+    scaler = torch.cuda.amp.grad_scaler.GradScaler()
+    amp_ctx = contextlib.nullcontext()
+    if args.fp16:
+        amp_ctx = torch.cuda.amp.autocast()
+
+    for batch_idx, (data, target) in enumerate(train_loader):
+        data, target = data.to(device), target.to(device)
+        optimizer.zero_grad()
+        with amp_ctx:
+            output = model(data)
+            print(output.dtype)
+
+
+def main():
+    # Training settings
+    parser = argparse.ArgumentParser(description="PyTorch MNIST Example")
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=1000,
+        metavar="N",
+        help="input batch size for training (default: 64)",
+    )
+    parser.add_argument(
+        "--test-batch-size",
+        type=int,
+        default=1000,
+        metavar="N",
+        help="input batch size for testing (default: 1000)",
+    )
+    parser.add_argument(
+        "--epochs",
+        type=int,
+        default=10,
+        metavar="N",
+        help="number of epochs to train (default: 14)",
+    )
+    parser.add_argument(
+        "--lr",
+        type=float,
+        default=1.0,
+        metavar="LR",
+        help="learning rate (default: 1.0)",
+    )
+    parser.add_argument(
+        "--gamma",
+        type=float,
+        default=0.7,
+        metavar="M",
+        help="Learning rate step gamma (default: 0.7)",
+    )
+    parser.add_argument(
+        "--no-cuda", action="store_true", default=False, help="disables CUDA training"
+    )
+    parser.add_argument(
+        "--seed", type=int, default=1, metavar="S", help="random seed (default: 1)"
+    )
+    parser.add_argument(
+        "--log-interval",
+        type=int,
+        default=10,
+        metavar="N",
+        help="how many batches to wait before logging training status",
+    )
+
+    parser.add_argument(
+        "--save-model",
+        action="store_true",
+        default=False,
+        help="For Saving the current Model",
+    )
+    parser.add_argument(
+        "--fp16",
+        action="store_true",
+        default=False,
+        help="For mixed precision training",
+    )
+
+    args = parser.parse_args()
+    use_cuda = not args.no_cuda and torch.cuda.is_available()
+
+    torch.manual_seed(args.seed)
+
+    device = torch.device("cuda" if use_cuda else "cpu")
+
+    kwargs = {"num_workers": 20, "pin_memory": True} if use_cuda else {}
+    train_loader = torch.utils.data.DataLoader(
+        datasets.MNIST(
+            "../data",
+            train=True,
+            download=True,
+            transform=transforms.Compose(
+                [
+                    transforms.ToTensor(),
+                    # here we remove norm to get sparse tensor with lots of zeros
+                    # transforms.Normalize((0.1307,), (0.3081,))
+                ]
+            ),
+        ),
+        batch_size=args.batch_size,
+        shuffle=True,
+        **kwargs,
+    )
+    test_loader = torch.utils.data.DataLoader(
+        datasets.MNIST(
+            "../data",
+            train=False,
+            transform=transforms.Compose(
+                [
+                    transforms.ToTensor(),
+                    # here we remove norm to get sparse tensor with lots of zeros
+                    # transforms.Normalize((0.1307,), (0.3081,))
+                ]
+            ),
+        ),
+        batch_size=args.test_batch_size,
+        shuffle=True,
+        **kwargs,
+    )
+
+    model = Net().to(device)
+    optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
+
+    scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
+
+    t0 = time.time()
+    for epoch in range(1, args.epochs + 1):
+        train(args, model, device, train_loader, optimizer, epoch)
+
+    torch.cuda.current_stream().synchronize()
+    t1 = time.time()
+
+    print(t1 - t0)
+
+    if args.save_model:
+        torch.save(model.state_dict(), "mnist_cnn.pt")
+
+
+if __name__ == "__main__":
+    main()