various

devel/export_onnx_normaldistribution.py

@@ -0,0 +1,59 @@
+import torch
+import torch.utils.data
+from torch import nn
+from torch.nn import functional as F
+
+from datetime import datetime
+
+now = datetime.now()
+
+currdatestring = now.strftime("%Y%m%d_%H%M%S")
+
+
+IN_DIMS = 28 * 28
+BATCH_SIZE = 10
+FEATURE_DIM = 20
+
+class VAE(nn.Module):
+    def __init__(self):
+        super(VAE, self).__init__()
+
+        self.fc1 = nn.Linear(784, 400)
+        self.fc21 = nn.Linear(400, FEATURE_DIM)
+        self.fc22 = nn.Linear(400, FEATURE_DIM)
+        self.fc3 = nn.Linear(FEATURE_DIM, 400)
+        self.fc4 = nn.Linear(400, 784)
+
+    def encode(self, x):
+        h1 = F.relu(self.fc1(x))
+        return self.fc21(h1), self.fc22(h1)
+
+    def reparameterize(self, mu, logvar):
+        std = torch.exp(0.5*logvar)
+
+        m = torch.distributions.normal.Normal(torch.tensor([0.0]), torch.tensor([1.0]))
+        eps = m.sample()
+#        m = torch.distributions.von_mises.VonMises(torch.tensor([1.0]), torch.tensor([1.0]))
+#        eps = m.sample()
+#        eps = torch.randn(BATCH_SIZE, FEATURE_DIM, device='cuda')
+        return eps #.mul(std).add_(mu)
+
+    def decode(self, z):
+        h3 = F.relu(self.fc3(z))
+        return torch.sigmoid(self.fc4(h3))
+
+    def forward(self, x):
+        mu, logvar = self.encode(x)
+        z = self.reparameterize(mu, logvar)
+        recon_x = self.decode(z)
+
+        return recon_x
+
+model = VAE()
+
+dummy_input = torch.randn(BATCH_SIZE, IN_DIMS, device='cpu')
+torch.onnx.export(model, dummy_input, "vae" + currdatestring + ".onnx", verbose=True)
+
+#torch.onnx.export(model, in_data, 'test.onnx', verbose=True)
+
+#torch.onnx.dynamo_export(model, dummy_in_data)

devel/export_onnx_power.py

@@ -0,0 +1,44 @@
+import torch
+import torchvision.transforms.functional as functional
+import os
+import sys
+
+script_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(os.path.join(script_dir, '../'))
+
+
+import data
+import models
+import torch.onnx
+
+os.environ["CUDA_VISIBLE_DEVICES"]=""
+
+print(torch.__version__)
+
+ckpt_path = '/local_data/fehlneas/Spherinator/lightning_logs/version_0/checkpoints/epoch=3-step=400.ckpt'
+
+
+model = models.RotationalVariationalAutoencoderPower() #.load_from_checkpoint(ckpt_path,map_location=torch.device('cpu'))
+model.eval()
+print(model)
+
+dummy_input = torch.randn(1,3,64,64)
+try:
+    torch.onnx.enable_log()
+    filepath="new.onnx"
+    model.to_onnx(filepath, export_params=True, opset_version=14,verbose=True)
+except Exception as inst:
+    print(type(inst))    # the exception type
+    print(inst.args)     # arguments stored in .args
+    print(inst) 
+
+print('############################')
+
+try:
+    torch.onnx.dynamo_export(model,dummy_input)
+except Exception as inst:
+    print(type(inst))    # the exception type
+    print(inst.args)     # arguments stored in .args
+    print(inst)
+
+

devel/export_onnx_svrae_norm.py

@@ -0,0 +1,44 @@
+import torch
+import torchvision.transforms.functional as functional
+import os
+import sys
+
+script_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(os.path.join(script_dir, '../'))
+
+
+import data
+import models
+import torch.onnx
+
+os.environ["CUDA_VISIBLE_DEVICES"]=""
+
+print(torch.__version__)
+
+ckpt_path = '/local_data/fehlneas/Spherinator/lightning_logs/version_0/checkpoints/epoch=3-step=400.ckpt'
+
+
+model = models.RotationalVariationalAutoencoderNorm.load_from_checkpoint(ckpt_path,map_location=torch.device('cpu'))
+model.eval()
+print(model)
+
+dummy_input = torch.randn(1,3,64,64)
+try:
+    torch.onnx.enable_log()
+    filepath="new.onnx"
+    model.to_onnx(filepath, export_params=True, opset_version=14,verbose=True)
+except Exception as inst:
+    print(type(inst))    # the exception type
+    print(inst.args)     # arguments stored in .args
+    print(inst) 
+
+print('############################')
+
+try:
+    torch.onnx.dynamo_export(model,dummy_input)
+except Exception as inst:
+    print(type(inst))    # the exception type
+    print(inst.args)     # arguments stored in .args
+    print(inst)
+
+

devel/export_onnx_vonmises.py

@@ -0,0 +1,29 @@
+import math
+import torch
+import torch.nn as nn
+from torch.distributions.utils import broadcast_all
+
+
+class Model(nn.Module):
+  def __init__(self):
+    super(Model, self).__init__()
+    self.mean = nn.Parameter(torch.zeros((2, 10)), requires_grad=False)
+    self.std = nn.Parameter(torch.ones((2, 10)), requires_grad=False)
+
+  def forward(self, x):
+    mean, std = broadcast_all(self.mean, self.std)
+    #m = torch.distributions.normal.Normal(torch.tensor([0.0]), torch.tensor([1.0]))
+    # mean, std = self.mean, self.std
+
+    m = torch.distributions.VonMises(torch.tensor([1.0]), torch.tensor([1.0]))
+    m.sample()
+    return m #-((x - mean) ** 2) / (2 * std ** 2) - math.log(math.sqrt(2 * math.pi))
+
+
+model = Model()
+in_data = torch.ones((2, 10))
+out_data = model(in_data)
+#torch.onnx.export(model, in_data, 'test.onnx', verbose=True)
+
+torch.onnx.dynamo_export(model, in_data)
+

devel/gaussian_vae.py

@@ -0,0 +1,142 @@
+from __future__ import print_function
+import argparse
+import torch
+import torch.utils.data
+from torch import nn, optim
+from torch.nn import functional as F
+from torchvision import datasets, transforms
+from torchvision.utils import save_image
+import torch.distributions as td
+
+
+parser = argparse.ArgumentParser(description='VAE MNIST Example')
+parser.add_argument('--batch-size', type=int, default=128, metavar='N',
+                    help='input batch size for training (default: 128)')
+parser.add_argument('--epochs', type=int, default=10, metavar='N',
+                    help='number of epochs to train (default: 10)')
+parser.add_argument('--no-cuda', action='store_true', default=False,
+                    help='enables 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')
+args = parser.parse_args()
+args.cuda = not args.no_cuda and torch.cuda.is_available()
+
+torch.manual_seed(args.seed)
+
+device = torch.device("cuda" if args.cuda else "cpu")
+
+kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}
+train_loader = torch.utils.data.DataLoader(
+    datasets.MNIST('../data', train=True, download=True,
+                   transform=transforms.ToTensor()),
+    batch_size=args.batch_size, shuffle=True, **kwargs)
+test_loader = torch.utils.data.DataLoader(
+    datasets.MNIST('../data', train=False, transform=transforms.ToTensor()),
+    batch_size=args.batch_size, shuffle=True, **kwargs)
+
+
+class VAE(nn.Module):
+    def __init__(self):
+        super(VAE, self).__init__()
+
+        self.fc1 = nn.Linear(784, 400)
+        self.fc21 = nn.Linear(400, 20)
+        self.fc22 = nn.Linear(400, 20)
+        self.fc3 = nn.Linear(20, 400)
+        self.fc4 = nn.Linear(400, 784)
+
+    def encode(self, x):
+        h1 = F.relu(self.fc1(x))
+        return self.fc21(h1), self.fc22(h1)
+
+    def decode(self, z):
+        h3 = F.relu(self.fc3(z))
+        # return torch.sigmoid(self.fc4(h3))
+        return self.fc4(h3)
+
+    def forward(self, x):
+        mu, logvar = self.encode(x.view(-1, 784))
+
+        std = logvar.exp().pow(0.5)         # logvar to std
+        q_z = td.normal.Normal(mu, std)     # create a torch distribution
+        z = q_z.rsample()                   # sample with reparameterization
+
+        return self.decode(z), q_z
+
+
+model = VAE().to(device)
+optimizer = optim.Adam(model.parameters(), lr=1e-3)
+
+
+# Reconstruction + KL divergence losses summed over all elements and batch
+def loss_function(recon_x, x, q_z):
+    BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduction='sum')
+    # You can also compute p(x|z) as below, for binary output it reduces
+    # to binary cross entropy error, for gaussian output it reduces to
+    # mean square error
+    # p_x = td.bernoulli.Bernoulli(logits=recon_x)
+    # BCE = -p_x.log_prob(x.view(-1, 784)).sum()
+
+    # see Appendix B from VAE paper:
+    # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
+    # https://arxiv.org/abs/1312.6114
+    # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
+
+    # KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
+    p_z = td.normal.Normal(torch.zeros_like(q_z.loc), torch.ones_like(q_z.scale))
+    KLD = td.kl_divergence(q_z, p_z).sum()
+
+    return BCE + KLD
+
+
+def train(epoch):
+    model.train()
+    train_loss = 0
+    for batch_idx, (data, _) in enumerate(train_loader):
+        data = data.to(device)
+        optimizer.zero_grad()
+        recon_batch, q_z = model(data)
+        loss = loss_function(recon_batch, data, q_z)
+        loss.backward()
+        train_loss += loss.item()
+        optimizer.step()
+        if batch_idx % args.log_interval == 0:
+            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
+                epoch, batch_idx * len(data), len(train_loader.dataset),
+                100. * batch_idx / len(train_loader),
+                loss.item() / len(data)))
+
+    print('====> Epoch: {} Average loss: {:.4f}'.format(
+          epoch, train_loss / len(train_loader.dataset)))
+
+
+def test(epoch):
+    model.eval()
+    test_loss = 0
+    with torch.no_grad():
+        for i, (data, _) in enumerate(test_loader):
+            data = data.to(device)
+            recon_batch, q_z = model(data)
+            test_loss += loss_function(recon_batch, data, q_z).item()
+            if i == 0:
+                n = min(data.size(0), 8)
+                comparison = torch.cat([data[:n],
+                                       recon_batch.view(args.batch_size, 1, 28, 28)[:n]])
+                save_image(comparison.cpu(),
+                           'results/reconstruction_' + str(epoch) + '.png', nrow=n)
+
+    test_loss /= len(test_loader.dataset)
+    print('====> Test set loss: {:.4f}'.format(test_loss))
+
+
+if __name__ == "__main__":
+    for epoch in range(1, args.epochs + 1):
+        train(epoch)
+        test(epoch)
+        with torch.no_grad():
+            sample = torch.randn(64, 20).to(device)
+            sample = model.decode(sample).cpu()
+            save_image(sample.view(64, 1, 28, 28),
+                       'results/sample_' + str(epoch) + '.png')

devel/onnx_export_example.py

@@ -0,0 +1,69 @@
+# Super Resolution model definition in PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+import os
+os.environ["CUDA_VISIBLE_DEVICES"]=""
+
+
+
+class SuperResolutionNet(nn.Module):
+    def __init__(self, upscale_factor, inplace=False):
+        super(SuperResolutionNet, self).__init__()
+
+        self.relu = nn.ReLU(inplace=inplace)
+        self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))
+        self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
+        self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
+        self.conv4 = nn.Conv2d(32, upscale_factor ** 2, (3, 3), (1, 1), (1, 1))
+        self.pixel_shuffle = nn.PixelShuffle(upscale_factor)
+
+        self._initialize_weights()
+
+    def forward(self, x):
+        x = self.relu(self.conv1(x))
+        x = self.relu(self.conv2(x))
+        x = self.relu(self.conv3(x))
+        x = self.pixel_shuffle(self.conv4(x))
+        x = torch.distributions.normal.Normal(x*0.1, x*0.5)
+        return x
+
+    def _initialize_weights(self):
+        init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))
+        init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))
+        init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))
+        init.orthogonal_(self.conv4.weight)
+
+# Create the super-resolution model by using the above model definition.
+torch_model = SuperResolutionNet(upscale_factor=3)
+
+
+# Load pretrained model weights
+model_url = 'https://s3.amazonaws.com/pytorch/test_data/export/superres_epoch100-44c6958e.pth'
+batch_size = 1    # just a random number
+
+# Initialize model with the pretrained weights
+map_location = lambda storage, loc: storage
+#if torch.cuda.is_available():
+#    map_location = None
+#torch_model.load_state_dict(model_zoo.load_url(model_url, map_location=map_location))
+
+# set the model to inference mode
+torch_model.eval()
+
+
+# Input to the model
+x = torch.randn(batch_size, 1, 224, 224, requires_grad=True)
+torch_out = torch_model(x)
+
+# Export the model
+torch.onnx.export(torch_model,               # model being run
+                  x,                         # model input (or a tuple for multiple inputs)
+                  "super_resolution.onnx",   # where to save the model (can be a file or file-like object)
+                  export_params=True,        # store the trained parameter weights inside the model file
+                  opset_version=10,          # the ONNX version to export the model to
+                  do_constant_folding=True,  # whether to execute constant folding for optimization
+                  input_names = ['input'],   # the model's input names
+                  output_names = ['output'], # the model's output names
+                  dynamic_axes={'input' : {0 : 'batch_size'},    # variable length axes
+                                'output' : {0 : 'batch_size'}})