import numpy as np
import os
import sys
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torch.nn.init as init
import torch.utils.data as data
import torch.utils.data.dataset as dataset
import torchvision.datasets as dset
import torchvision.transforms as transforms
from torch.autograd import Variable
import torchvision.utils as v_utils
import matplotlib.pyplot as plt
import cv2
import math
from collections import OrderedDict
import copy
import time
from model.utils import DataLoader
from model.final_future_prediction_with_memory_spatial_sumonly_weight_ranking_top1 import *
from sklearn.metrics import roc_auc_score
from utils import *
import random

import argparse

if __name__ == '__main__':

    parser = argparse.ArgumentParser(description="MNAD")
    parser.add_argument('--gpus', nargs='+', type=str, help='gpus')
    parser.add_argument('--batch_size', type=int, default=2, help='batch size for training')
    parser.add_argument('--test_batch_size', type=int, default=1, help='batch size for test')
    parser.add_argument('--epochs', type=int, default=40, help='number of epochs for training')
    parser.add_argument('--loss_compact', type=float, default=0.1, help='weight of the feature compactness loss')
    parser.add_argument('--loss_separate', type=float, default=0.1, help='weight of the feature separateness loss')
    parser.add_argument('--h', type=int, default=256, help='height of input images')
    parser.add_argument('--w', type=int, default=256, help='width of input images')
    parser.add_argument('--c', type=int, default=3, help='channel of input images')
    parser.add_argument('--lr', type=float, default=2e-4, help='initial learning rate')
    parser.add_argument('--t_length', type=int, default=5, help='length of the frame sequences')
    parser.add_argument('--fdim', type=int, default=512, help='channel dimension of the features')
    parser.add_argument('--mdim', type=int, default=512, help='channel dimension of the memory items')
    parser.add_argument('--msize', type=int, default=10, help='number of the memory items')
    parser.add_argument('--alpha', type=float, default=0.6, help='weight for the anomality score')
    parser.add_argument('--num_workers', type=int, default=2, help='number of workers for the train loader')
    parser.add_argument('--num_workers_test', type=int, default=1, help='number of workers for the test loader')
    parser.add_argument('--dataset_type', type=str, default='ped2', help='type of dataset: ped2, avenue, shanghai')
    parser.add_argument('--dataset_path', type=str, default='./dataset/', help='directory of data')
    parser.add_argument('--exp_dir', type=str, default='log_position', help='directory of log')
    parser.add_argument('--model_dir', type=str, help='directory of model')
    parser.add_argument('--m_items_dir', type=str, help='directory of model')

    args = parser.parse_args()
    #my code: load model
    # exp_path = './exp/'
    # model_dir = 'model.pth'
    # m_items_dir = 'keys.pt'
    # # args.model_dir = 'Ped2_prediction_model.pth'
    # # args.m_items_dir = 'Ped2_prediction_keys.pt'
    # args.model_dir = exp_path + args.dataset_type + '/log/' + model_dir
    # args.m_items_dir = exp_path + args.dataset_type + '/log/' + m_items_dir

    torch.manual_seed(2020)

    os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
    if args.gpus is None:
        gpus = "0"
        os.environ["CUDA_VISIBLE_DEVICES"]= gpus
    else:
        gpus = ""
        for i in range(len(args.gpus)):
            gpus = gpus + args.gpus[i] + ","
        os.environ["CUDA_VISIBLE_DEVICES"]= gpus[:-1]

    torch.backends.cudnn.enabled = True # make sure to use cudnn for computational performance

    train_folder = args.dataset_path+args.dataset_type+"/training/frames"
    test_folder = args.dataset_path+args.dataset_type+"/testing/frames"

    # Loading dataset
    train_dataset = DataLoader(train_folder, transforms.Compose([
                 transforms.ToTensor(),
                 ]), resize_height=args.h, resize_width=args.w, time_step=args.t_length-1)

    test_dataset = DataLoader(test_folder, transforms.Compose([
                 transforms.ToTensor(),
                 ]), resize_height=args.h, resize_width=args.w, time_step=args.t_length-1)

    train_size = len(train_dataset)
    test_size = len(test_dataset)

    train_batch = data.DataLoader(train_dataset, batch_size = args.batch_size,
                                  shuffle=True, num_workers=args.num_workers, drop_last=True)
    test_batch = data.DataLoader(test_dataset, batch_size = args.test_batch_size,
                                 shuffle=False, num_workers=args.num_workers_test, drop_last=False)


    # Model setting
    model = convAE(args.c, args.t_length, args.msize, args.fdim, args.mdim)
    params_encoder =  list(model.encoder.parameters())
    params_decoder = list(model.decoder.parameters())
    params = params_encoder + params_decoder
    optimizer = torch.optim.Adam(params, lr = args.lr)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer,T_max =args.epochs)
    # model = torch.load(args.model_dir)
    model.cuda()
    # m_items = torch.load(args.m_items_dir)


    # Report the training process
    log_dir = os.path.join('./exp', args.dataset_type, args.exp_dir)
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
    orig_stdout = sys.stdout
    f = open(os.path.join(log_dir, 'log.txt'),'w')
    sys.stdout= f

    loss_func_mse = nn.MSELoss(reduction='none')

    # Training

    m_items = F.normalize(torch.rand((args.msize, args.mdim), dtype=torch.float), dim=1).cuda() # Initialize the memory items
    m_dis = F.normalize(torch.zeros((32, 32, 1), dtype=torch.float), dim=-1).cuda()
    # timing_signal = get_timing_signal_1d(args.h, args.w)
    # timing_signal = timing_signal.repeat(args.t_length-1,1,1)

    get_memory_dis = False
    for epoch in range(args.epochs):
        if epoch == args.epochs-1:
            get_memory_dis = True

        labels_list = []
        model.train()

        start = time.time()
        for j,(imgs) in enumerate(train_batch):

            # imgs[:, 0:12] = imgs[:, 0:12] + timing_signal
            imgs = Variable(imgs).cuda()

            outputs, _, _, m_items, softmax_score_query, softmax_score_memory, separateness_loss, compactness_loss, m_dis = model.forward(imgs[:,0:12], m_items, m_dis, True, get_memory_dis)


            optimizer.zero_grad()
            loss_pixel = torch.mean(loss_func_mse(outputs, imgs[:,12:]))
            loss = loss_pixel + args.loss_compact * compactness_loss + args.loss_separate * separateness_loss
            loss.backward(retain_graph=True)
            optimizer.step()

        scheduler.step()

        print('----------------------------------------')
        print('Epoch:', epoch+1)
        print('Loss: Reconstruction {:.6f}/ Compactness {:.6f}/ Separateness {:.6f}'.format(loss_pixel.item(), compactness_loss.item(), separateness_loss.item()))
        print('Memory_items:')
        print(m_items)
        print('----------------------------------------')

    print('Training is finished')
    # Save the model and the memory items
    torch.save(model, os.path.join(log_dir, 'model.pth'))
    torch.save(m_items, os.path.join(log_dir, 'keys.pt'))
    torch.save(m_dis, os.path.join(log_dir, 'dis.pt'))

    sys.stdout = orig_stdout
    f.close()




4d14a93ebf8f93f3d9e682643700e75dd824f307
<mask token>

<mask token>
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='MNAD')
    parser.add_argument('--gpus', nargs='+', type=str, help='gpus')
    parser.add_argument('--batch_size', type=int, default=2, help=
        'batch size for training')
    parser.add_argument('--test_batch_size', type=int, default=1, help=
        'batch size for test')
    parser.add_argument('--epochs', type=int, default=40, help=
        'number of epochs for training')
    parser.add_argument('--loss_compact', type=float, default=0.1, help=
        'weight of the feature compactness loss')
    parser.add_argument('--loss_separate', type=float, default=0.1, help=
        'weight of the feature separateness loss')
    parser.add_argument('--h', type=int, default=256, help=
        'height of input images')
    parser.add_argument('--w', type=int, default=256, help=
        'width of input images')
    parser.add_argument('--c', type=int, default=3, help=
        'channel of input images')
    parser.add_argument('--lr', type=float, default=0.0002, help=
        'initial learning rate')
    parser.add_argument('--t_length', type=int, default=5, help=
        'length of the frame sequences')
    parser.add_argument('--fdim', type=int, default=512, help=
        'channel dimension of the features')
    parser.add_argument('--mdim', type=int, default=512, help=
        'channel dimension of the memory items')
    parser.add_argument('--msize', type=int, default=10, help=
        'number of the memory items')
    parser.add_argument('--alpha', type=float, default=0.6, help=
        'weight for the anomality score')
    parser.add_argument('--num_workers', type=int, default=2, help=
        'number of workers for the train loader')
    parser.add_argument('--num_workers_test', type=int, default=1, help=
        'number of workers for the test loader')
    parser.add_argument('--dataset_type', type=str, default='ped2', help=
        'type of dataset: ped2, avenue, shanghai')
    parser.add_argument('--dataset_path', type=str, default='./dataset/',
        help='directory of data')
    parser.add_argument('--exp_dir', type=str, default='log_position', help
        ='directory of log')
    parser.add_argument('--model_dir', type=str, help='directory of model')
    parser.add_argument('--m_items_dir', type=str, help='directory of model')
    args = parser.parse_args()
    torch.manual_seed(2020)
    os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
    if args.gpus is None:
        gpus = '0'
        os.environ['CUDA_VISIBLE_DEVICES'] = gpus
    else:
        gpus = ''
        for i in range(len(args.gpus)):
            gpus = gpus + args.gpus[i] + ','
        os.environ['CUDA_VISIBLE_DEVICES'] = gpus[:-1]
    torch.backends.cudnn.enabled = True
    train_folder = args.dataset_path + args.dataset_type + '/training/frames'
    test_folder = args.dataset_path + args.dataset_type + '/testing/frames'
    train_dataset = DataLoader(train_folder, transforms.Compose([transforms
        .ToTensor()]), resize_height=args.h, resize_width=args.w, time_step
        =args.t_length - 1)
    test_dataset = DataLoader(test_folder, transforms.Compose([transforms.
        ToTensor()]), resize_height=args.h, resize_width=args.w, time_step=
        args.t_length - 1)
    train_size = len(train_dataset)
    test_size = len(test_dataset)
    train_batch = data.DataLoader(train_dataset, batch_size=args.batch_size,
        shuffle=True, num_workers=args.num_workers, drop_last=True)
    test_batch = data.DataLoader(test_dataset, batch_size=args.
        test_batch_size, shuffle=False, num_workers=args.num_workers_test,
        drop_last=False)
    model = convAE(args.c, args.t_length, args.msize, args.fdim, args.mdim)
    params_encoder = list(model.encoder.parameters())
    params_decoder = list(model.decoder.parameters())
    params = params_encoder + params_decoder
    optimizer = torch.optim.Adam(params, lr=args.lr)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.
        epochs)
    model.cuda()
    log_dir = os.path.join('./exp', args.dataset_type, args.exp_dir)
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
    orig_stdout = sys.stdout
    f = open(os.path.join(log_dir, 'log.txt'), 'w')
    sys.stdout = f
    loss_func_mse = nn.MSELoss(reduction='none')
    m_items = F.normalize(torch.rand((args.msize, args.mdim), dtype=torch.
        float), dim=1).cuda()
    m_dis = F.normalize(torch.zeros((32, 32, 1), dtype=torch.float), dim=-1
        ).cuda()
    get_memory_dis = False
    for epoch in range(args.epochs):
        if epoch == args.epochs - 1:
            get_memory_dis = True
        labels_list = []
        model.train()
        start = time.time()
        for j, imgs in enumerate(train_batch):
            imgs = Variable(imgs).cuda()
            (outputs, _, _, m_items, softmax_score_query,
                softmax_score_memory, separateness_loss, compactness_loss,
                m_dis) = (model.forward(imgs[:, 0:12], m_items, m_dis, True,
                get_memory_dis))
            optimizer.zero_grad()
            loss_pixel = torch.mean(loss_func_mse(outputs, imgs[:, 12:]))
            loss = (loss_pixel + args.loss_compact * compactness_loss + 
                args.loss_separate * separateness_loss)
            loss.backward(retain_graph=True)
            optimizer.step()
        scheduler.step()
        print('----------------------------------------')
        print('Epoch:', epoch + 1)
        print(
            'Loss: Reconstruction {:.6f}/ Compactness {:.6f}/ Separateness {:.6f}'
            .format(loss_pixel.item(), compactness_loss.item(),
            separateness_loss.item()))
        print('Memory_items:')
        print(m_items)
        print('----------------------------------------')
    print('Training is finished')
    torch.save(model, os.path.join(log_dir, 'model.pth'))
    torch.save(m_items, os.path.join(log_dir, 'keys.pt'))
    torch.save(m_dis, os.path.join(log_dir, 'dis.pt'))
    sys.stdout = orig_stdout
    f.close()

import numpy as np
import os
import sys
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torch.nn.init as init
import torch.utils.data as data
import torch.utils.data.dataset as dataset
import torchvision.datasets as dset
import torchvision.transforms as transforms
from torch.autograd import Variable
import torchvision.utils as v_utils
import matplotlib.pyplot as plt
import cv2
import math
from collections import OrderedDict
import copy
import time
from model.utils import DataLoader
from model.final_future_prediction_with_memory_spatial_sumonly_weight_ranking_top1 import *
from sklearn.metrics import roc_auc_score
from utils import *
import random
import argparse
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='MNAD')
    parser.add_argument('--gpus', nargs='+', type=str, help='gpus')
    parser.add_argument('--batch_size', type=int, default=2, help=
        'batch size for training')
    parser.add_argument('--test_batch_size', type=int, default=1, help=
        'batch size for test')
    parser.add_argument('--epochs', type=int, default=40, help=
        'number of epochs for training')
    parser.add_argument('--loss_compact', type=float, default=0.1, help=
        'weight of the feature compactness loss')
    parser.add_argument('--loss_separate', type=float, default=0.1, help=
        'weight of the feature separateness loss')
    parser.add_argument('--h', type=int, default=256, help=
        'height of input images')
    parser.add_argument('--w', type=int, default=256, help=
        'width of input images')
    parser.add_argument('--c', type=int, default=3, help=
        'channel of input images')
    parser.add_argument('--lr', type=float, default=0.0002, help=
        'initial learning rate')
    parser.add_argument('--t_length', type=int, default=5, help=
        'length of the frame sequences')
    parser.add_argument('--fdim', type=int, default=512, help=
        'channel dimension of the features')
    parser.add_argument('--mdim', type=int, default=512, help=
        'channel dimension of the memory items')
    parser.add_argument('--msize', type=int, default=10, help=
        'number of the memory items')
    parser.add_argument('--alpha', type=float, default=0.6, help=
        'weight for the anomality score')
    parser.add_argument('--num_workers', type=int, default=2, help=
        'number of workers for the train loader')
    parser.add_argument('--num_workers_test', type=int, default=1, help=
        'number of workers for the test loader')
    parser.add_argument('--dataset_type', type=str, default='ped2', help=
        'type of dataset: ped2, avenue, shanghai')
    parser.add_argument('--dataset_path', type=str, default='./dataset/',
        help='directory of data')
    parser.add_argument('--exp_dir', type=str, default='log_position', help
        ='directory of log')
    parser.add_argument('--model_dir', type=str, help='directory of model')
    parser.add_argument('--m_items_dir', type=str, help='directory of model')
    args = parser.parse_args()
    torch.manual_seed(2020)
    os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
    if args.gpus is None:
        gpus = '0'
        os.environ['CUDA_VISIBLE_DEVICES'] = gpus
    else:
        gpus = ''
        for i in range(len(args.gpus)):
            gpus = gpus + args.gpus[i] + ','
        os.environ['CUDA_VISIBLE_DEVICES'] = gpus[:-1]
    torch.backends.cudnn.enabled = True
    train_folder = args.dataset_path + args.dataset_type + '/training/frames'
    test_folder = args.dataset_path + args.dataset_type + '/testing/frames'
    train_dataset = DataLoader(train_folder, transforms.Compose([transforms
        .ToTensor()]), resize_height=args.h, resize_width=args.w, time_step
        =args.t_length - 1)
    test_dataset = DataLoader(test_folder, transforms.Compose([transforms.
        ToTensor()]), resize_height=args.h, resize_width=args.w, time_step=
        args.t_length - 1)
    train_size = len(train_dataset)
    test_size = len(test_dataset)
    train_batch = data.DataLoader(train_dataset, batch_size=args.batch_size,
        shuffle=True, num_workers=args.num_workers, drop_last=True)
    test_batch = data.DataLoader(test_dataset, batch_size=args.
        test_batch_size, shuffle=False, num_workers=args.num_workers_test,
        drop_last=False)
    model = convAE(args.c, args.t_length, args.msize, args.fdim, args.mdim)
    params_encoder = list(model.encoder.parameters())
    params_decoder = list(model.decoder.parameters())
    params = params_encoder + params_decoder
    optimizer = torch.optim.Adam(params, lr=args.lr)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.
        epochs)
    model.cuda()
    log_dir = os.path.join('./exp', args.dataset_type, args.exp_dir)
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
    orig_stdout = sys.stdout
    f = open(os.path.join(log_dir, 'log.txt'), 'w')
    sys.stdout = f
    loss_func_mse = nn.MSELoss(reduction='none')
    m_items = F.normalize(torch.rand((args.msize, args.mdim), dtype=torch.
        float), dim=1).cuda()
    m_dis = F.normalize(torch.zeros((32, 32, 1), dtype=torch.float), dim=-1
        ).cuda()
    get_memory_dis = False
    for epoch in range(args.epochs):
        if epoch == args.epochs - 1:
            get_memory_dis = True
        labels_list = []
        model.train()
        start = time.time()
        for j, imgs in enumerate(train_batch):
            imgs = Variable(imgs).cuda()
            (outputs, _, _, m_items, softmax_score_query,
                softmax_score_memory, separateness_loss, compactness_loss,
                m_dis) = (model.forward(imgs[:, 0:12], m_items, m_dis, True,
                get_memory_dis))
            optimizer.zero_grad()
            loss_pixel = torch.mean(loss_func_mse(outputs, imgs[:, 12:]))
            loss = (loss_pixel + args.loss_compact * compactness_loss + 
                args.loss_separate * separateness_loss)
            loss.backward(retain_graph=True)
            optimizer.step()
        scheduler.step()
        print('----------------------------------------')
        print('Epoch:', epoch + 1)
        print(
            'Loss: Reconstruction {:.6f}/ Compactness {:.6f}/ Separateness {:.6f}'
            .format(loss_pixel.item(), compactness_loss.item(),
            separateness_loss.item()))
        print('Memory_items:')
        print(m_items)
        print('----------------------------------------')
    print('Training is finished')
    torch.save(model, os.path.join(log_dir, 'model.pth'))
    torch.save(m_items, os.path.join(log_dir, 'keys.pt'))
    torch.save(m_dis, os.path.join(log_dir, 'dis.pt'))
    sys.stdout = orig_stdout
    f.close()

import numpy as np
import os
import sys
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torch.nn.init as init
import torch.utils.data as data
import torch.utils.data.dataset as dataset
import torchvision.datasets as dset
import torchvision.transforms as transforms
from torch.autograd import Variable
import torchvision.utils as v_utils
import matplotlib.pyplot as plt
import cv2
import math
from collections import OrderedDict
import copy
import time
from model.utils import DataLoader
from model.final_future_prediction_with_memory_spatial_sumonly_weight_ranking_top1 import *
from sklearn.metrics import roc_auc_score
from utils import *
import random

import argparse

if __name__ == '__main__':

    parser = argparse.ArgumentParser(description="MNAD")
    parser.add_argument('--gpus', nargs='+', type=str, help='gpus')
    parser.add_argument('--batch_size', type=int, default=2, help='batch size for training')
    parser.add_argument('--test_batch_size', type=int, default=1, help='batch size for test')
    parser.add_argument('--epochs', type=int, default=40, help='number of epochs for training')
    parser.add_argument('--loss_compact', type=float, default=0.1, help='weight of the feature compactness loss')
    parser.add_argument('--loss_separate', type=float, default=0.1, help='weight of the feature separateness loss')
    parser.add_argument('--h', type=int, default=256, help='height of input images')
    parser.add_argument('--w', type=int, default=256, help='width of input images')
    parser.add_argument('--c', type=int, default=3, help='channel of input images')
    parser.add_argument('--lr', type=float, default=2e-4, help='initial learning rate')
    parser.add_argument('--t_length', type=int, default=5, help='length of the frame sequences')
    parser.add_argument('--fdim', type=int, default=512, help='channel dimension of the features')
    parser.add_argument('--mdim', type=int, default=512, help='channel dimension of the memory items')
    parser.add_argument('--msize', type=int, default=10, help='number of the memory items')
    parser.add_argument('--alpha', type=float, default=0.6, help='weight for the anomality score')
    parser.add_argument('--num_workers', type=int, default=2, help='number of workers for the train loader')
    parser.add_argument('--num_workers_test', type=int, default=1, help='number of workers for the test loader')
    parser.add_argument('--dataset_type', type=str, default='ped2', help='type of dataset: ped2, avenue, shanghai')
    parser.add_argument('--dataset_path', type=str, default='./dataset/', help='directory of data')
    parser.add_argument('--exp_dir', type=str, default='log_position', help='directory of log')
    parser.add_argument('--model_dir', type=str, help='directory of model')
    parser.add_argument('--m_items_dir', type=str, help='directory of model')

    args = parser.parse_args()
    #my code: load model
    # exp_path = './exp/'
    # model_dir = 'model.pth'
    # m_items_dir = 'keys.pt'
    # # args.model_dir = 'Ped2_prediction_model.pth'
    # # args.m_items_dir = 'Ped2_prediction_keys.pt'
    # args.model_dir = exp_path + args.dataset_type + '/log/' + model_dir
    # args.m_items_dir = exp_path + args.dataset_type + '/log/' + m_items_dir

    torch.manual_seed(2020)

    os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
    if args.gpus is None:
        gpus = "0"
        os.environ["CUDA_VISIBLE_DEVICES"]= gpus
    else:
        gpus = ""
        for i in range(len(args.gpus)):
            gpus = gpus + args.gpus[i] + ","
        os.environ["CUDA_VISIBLE_DEVICES"]= gpus[:-1]

    torch.backends.cudnn.enabled = True # make sure to use cudnn for computational performance

    train_folder = args.dataset_path+args.dataset_type+"/training/frames"
    test_folder = args.dataset_path+args.dataset_type+"/testing/frames"

    # Loading dataset
    train_dataset = DataLoader(train_folder, transforms.Compose([
                 transforms.ToTensor(),
                 ]), resize_height=args.h, resize_width=args.w, time_step=args.t_length-1)

    test_dataset = DataLoader(test_folder, transforms.Compose([
                 transforms.ToTensor(),
                 ]), resize_height=args.h, resize_width=args.w, time_step=args.t_length-1)

    train_size = len(train_dataset)
    test_size = len(test_dataset)

    train_batch = data.DataLoader(train_dataset, batch_size = args.batch_size,
                                  shuffle=True, num_workers=args.num_workers, drop_last=True)
    test_batch = data.DataLoader(test_dataset, batch_size = args.test_batch_size,
                                 shuffle=False, num_workers=args.num_workers_test, drop_last=False)


    # Model setting
    model = convAE(args.c, args.t_length, args.msize, args.fdim, args.mdim)
    params_encoder =  list(model.encoder.parameters())
    params_decoder = list(model.decoder.parameters())
    params = params_encoder + params_decoder
    optimizer = torch.optim.Adam(params, lr = args.lr)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer,T_max =args.epochs)
    # model = torch.load(args.model_dir)
    model.cuda()
    # m_items = torch.load(args.m_items_dir)


    # Report the training process
    log_dir = os.path.join('./exp', args.dataset_type, args.exp_dir)
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
    orig_stdout = sys.stdout
    f = open(os.path.join(log_dir, 'log.txt'),'w')
    sys.stdout= f

    loss_func_mse = nn.MSELoss(reduction='none')

    # Training

    m_items = F.normalize(torch.rand((args.msize, args.mdim), dtype=torch.float), dim=1).cuda() # Initialize the memory items
    m_dis = F.normalize(torch.zeros((32, 32, 1), dtype=torch.float), dim=-1).cuda()
    # timing_signal = get_timing_signal_1d(args.h, args.w)
    # timing_signal = timing_signal.repeat(args.t_length-1,1,1)

    get_memory_dis = False
    for epoch in range(args.epochs):
        if epoch == args.epochs-1:
            get_memory_dis = True

        labels_list = []
        model.train()

        start = time.time()
        for j,(imgs) in enumerate(train_batch):

            # imgs[:, 0:12] = imgs[:, 0:12] + timing_signal
            imgs = Variable(imgs).cuda()

            outputs, _, _, m_items, softmax_score_query, softmax_score_memory, separateness_loss, compactness_loss, m_dis = model.forward(imgs[:,0:12], m_items, m_dis, True, get_memory_dis)


            optimizer.zero_grad()
            loss_pixel = torch.mean(loss_func_mse(outputs, imgs[:,12:]))
            loss = loss_pixel + args.loss_compact * compactness_loss + args.loss_separate * separateness_loss
            loss.backward(retain_graph=True)
            optimizer.step()

        scheduler.step()

        print('----------------------------------------')
        print('Epoch:', epoch+1)
        print('Loss: Reconstruction {:.6f}/ Compactness {:.6f}/ Separateness {:.6f}'.format(loss_pixel.item(), compactness_loss.item(), separateness_loss.item()))
        print('Memory_items:')
        print(m_items)
        print('----------------------------------------')

    print('Training is finished')
    # Save the model and the memory items
    torch.save(model, os.path.join(log_dir, 'model.pth'))
    torch.save(m_items, os.path.join(log_dir, 'keys.pt'))
    torch.save(m_dis, os.path.join(log_dir, 'dis.pt'))

    sys.stdout = orig_stdout
    f.close()



