/conv-tt-lstm

Primary LanguagePythonOtherNOASSERTION

Convolutional Tensor-Train LSTM (Conv-TT-LSTM)

Intro

PyTorch implementations of the paper, 'Convolutional Tensor-Train LSTM for Spatio-Temporal Learning', NeurIPS 2020. [project page]

  • code/ (original): The original implementation of the paper.
  • code_opt/ (optimized): The optimized implementation to accelerate training.
    • The details of optimization tricks are presented at ECCV 2020 tutorial, 'Mixed Precision Training for Convolutional Tensor-Train LSTM' [slides] [video]

License

Copyright (c) 2020 NVIDIA Corporation. All rights reserved. This work is licensed under a NVIDIA Open Source Non-commercial license.

Dataset

  • Moving-MNIST-2

    • Generator [link]
    • Save the output of generate_moving_mnist(...) in npz format:
    np.savez(moving-mnist-train.npz, data)
    
  • KTH action

Evaluation of multi-steps prediction

Higher PSNR/SSIM and lower MSE/LPIPS values indicate better predictive results. # of FLOPs denotes the multiplications for one-step prediction per sample, and Time(m) represents the clock time (in minutes) required by training the model for one epoch (10,000 samples)

  • Moving-MNIST-2 dataset

  • KTH action dataset

Contacts

This code was written by Wonmin Byeon (wbyeon@nvidia.com) and Jiahao Su (jiahaosu@terpmail.umd.edu).