6.get_pytorch_IR.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



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)

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([1, 29, 29])
random_target = torch.randn([10])
#scripted_module = torch.jit.trace(my_module, random_input)


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


lr = 0.01
optimizer = torchopt.sgd(lr)



#opt_state = optimizer.init(params)

def train_one_batch(params, imgs, targets, opt_state):
    grads, (loss, preds) = grad_and_preds_and_loss(params, imgs, targets)
    updates, opt_state = optimizer.update(grads, opt_state)
    params = torchopt.apply_updates(params, updates)
    return params, loss, preds

for imgs, labels in trainset_dataloader:

    def label_to_target(label):
        ret = [0] * 10
        ret[label] = 1
        return ret

    targets = torch.Tensor([label_to_target(label) for label in labels.tolist()]).float()
    imgs, targets = (imgs.to(device), targets.to(device))


    # 1. functorch

    fmodule, params, buffers = functorch.make_functional_with_buffers(my_module)
    print(type(params))
    print(type(buffers))

    from torch._subclasses.fake_tensor import FakeTensorMode
    fake_mode = FakeTensorMode()

    def to_fake_tensor(fake_mode, tensor):
        if isinstance(tensor, torch.Tensor):
            return fake_mode.from_tensor(tensor)
        elif isinstance(tensor, tuple):
            rets = tuple()
            for elem in tensor:
                ret = to_fake_tensor(fake_mode, elem)
                rets = rets + (ret, )
            return rets

        assert False

    def print_compile_fn(fx_module, args):
        fx_module.print_readable()
        return fx_module

    def compute_loss(prarams, buffers, datas, labels):
        preds = fmodule(params, buffers, datas)
        loss = loss_fn(preds, labels)
        return loss

    from functorch.compile import aot_function, compiled_function
    from torch._guards import detect_fake_mode
    aot_fn = aot_function(compute_loss, fw_compiler = print_compile_fn, bw_compiler = print_compile_fn)


    for elem in params:
        print(type(elem))

    def clone_tensor(tensor, require_grad):
        if isinstance(tensor, torch.Tensor):
            return tensor.clone().detach().requires_grad_(require_grad)
        elif isinstance(tensor, tuple):
            rets = tuple()
            for elem in tensor:
                ret = clone_tensor(elem, require_grad)
                rets = rets + (ret, )
            return rets

    params_clone = clone_tensor(params, True)
    buffers_clone = clone_tensor(buffers, True)
    imgs_clone = clone_tensor(imgs, True)
    targets_clone = clone_tensor(targets, True)
    the_loss = aot_fn(params_clone, buffers_clone, imgs_clone, targets_clone)
    the_loss.backward()

    exit()