inference.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.parameter import Parameter
from torch.nn.modules.module import Module

import math
import numpy as np
import pandas as pd

import sys
sys.path.append('./lib/')
from pkl_process import *
from utils import load_graphdata_channel_my, compute_val_loss_sttn, image_to_patches
import tools

from time import time
import shutil
import argparse
import configparser
from tensorboardX import SummaryWriter
import os

from ST_Transformer_new import STTransformer # STTN model with linear layer to get positional embedding
from ST_Transformer_new_sinembedding import STTransformer_sinembedding #STTN model with sin()/cos() to get positional embedding, the same as "Attention is all your need"
from VQ_VAE import VQVAE
#%%

if __name__ == '__main__':
    
    
    params_path = './Experiment/debug/' ## Path for saving network parameters
    print('params_path:', params_path)
    # filename = './PEMSD7/V_25_r1_d0_w0_astcgn.npz' ## Data generated by prepareData.py
    # num_of_hours, num_of_days, num_of_weeks = 1, 0, 0 ## The same setting as prepareData.py

    param_file = 'epoch_248000.params'
    
    ### Training Hyparameter
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    DEVICE = device
    batch_size = 36
    batch_length = 18
    learning_rate = 0.01
    epochs = 1000000
    
    ### Generate Data Loader
    # train_loader, train_target_tensor, val_loader, val_target_tensor, test_loader, test_target_tensor, _mean, _std = load_graphdata_channel_my(
    #     filename, num_of_hours, num_of_days, num_of_weeks, DEVICE, batch_size)
    # import ipdb; ipdb.set_trace()

    dataset_path = '/media/ytzheng/3EA48EC9A48E835F/CARLA_DATA/Town01_copycat'
    eval_episodes = tools.load_episodes(dataset_path)
    generator = tools.sample_episodes(
        eval_episodes, batch_length
    )
    eval_dataset = tools.from_generator(generator, batch_size)
    
    ### Construct Network
    net =  VQVAE()
    net.to(device)
    print(net)
    param_path = os.path.join(params_path, param_file)
    print("Loading params from: ", param_path)
    net.load_state_dict(torch.load(param_path))

    ### Inference
    num_episodes = 1000
    for episode in range(num_episodes):
        net.eval()
        eval_batch = next(eval_dataset)
        patch_images = image_to_patches(torch.Tensor(eval_batch['image']).to(device))
        a_q, _ = net(patch_images.permute(0, 2, 1, 3))
        import ipdb; ipdb.set_trace()