8.py

#!/bin/env python
import time
from module import *
from typing import List
import torch

import torch.nn as nn
import torch.nn.functional as F

from torchvision.datasets import MNIST
from torch.utils.data import DataLoader
from torchvision import transforms as tfs
import torchopt
import functorch

import aot

device = torch.device("cpu")

class MyTransform:
    def __init__(self):
        pass

    def __call__(self, img):
        img = tfs.functional.rotate(img, -90)
        img = tfs.functional.hflip(img)
        img = img.map_(img, lambda a, b: 1.0 - a)
        return img

trainset = MNIST(root = "",
                #split = "digits",
                train = True,
                download = True,
                transform = tfs.Compose([tfs.PILToTensor(),
                        tfs.Resize(29),
                        tfs.ConvertImageDtype(torch.float32),
                        MyTransform()
                        ]))
testset = MNIST(root = "",
                #split = "digits",
                train = False,
                download = True,
                transform = tfs.Compose([tfs.PILToTensor(),
                        tfs.Resize(29),
                        tfs.ConvertImageDtype(torch.float32),
                        MyTransform()
                        ]))

trainset_size = len(trainset)
testset_size = len(testset)

batch_size = 16

trainset_dataloader = DataLoader(trainset, batch_size = 16, shuffle = False);
testset_dataloader = DataLoader(testset, batch_size = 16, shuffle = False);



my_module = MyModule()

random_input = torch.randn([batch_size, 1, 29, 29])
random_target = torch.randn([16, 10], requires_grad = False)
preds = my_module(random_input)

loss_fn = nn.CrossEntropyLoss()
loss_fn.to(device)

import my_graph_capture

ts_mod, states = my_graph_capture.capture_TorchScript(my_module, 0, random_input)

print(ts_mod)
print(states)