main_natural_mov.py

'''
Author: GhMa
Date: 2023-03-10 15:30:43
LastEditors: error: error: git config user.name & please set dead value or install git && error: git config user.email & please set dead value or install git & please set dead value or install git
LastEditTime: 2023-04-27 22:16:40
'''
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
from torch.utils.tensorboard import SummaryWriter

import os
import uuid
import argparse
import glob
import shutil
import importlib

from models import *
from crits import *
from utils import progress_bar, visualize_grad_norms, visualize_weight_norms
from utils import Monitor, visualize_spiking_rates
from dataset_utils import prepare_natural_movie

import torchvision
import torchvision.transforms as transforms


os.environ["CUDA_VISIBLE_DEVICES"] = "0"
COLOR_PURPLE = "\033[1;35m"    
COLOR_CYAN   = "\033[1;36m"     
COLOR_END    = '\033[0m'

parser = argparse.ArgumentParser(description='SNN Exps.')

parser.add_argument('--lr', default=0.0005, type=float, help='learning rate')
parser.add_argument('--epochs', default=64, type=int, help='#epoch')
parser.add_argument('--minibatch', default=64, type=int, 
                    help='mini-batch size')
# FOR NEURON
parser.add_argument('--neuron', default='LIF', type=str, help='neuron type')
parser.add_argument('--threshold', default=1.0, type=float, 
                    help='spiking thresh')
parser.add_argument(
    '--tau', default=2.0, type=float, 
    help='initial tau, membrane potential decay constant'
)
parser.add_argument(
    '--sigma', default=0.2, type=float, 
    help='std of p_epsilon, specifically for Noisy LIF'
)

# FOR NETWORK ARCH.
parser.add_argument(
    '--x_dim', default=90, type=int, 
    help='input width. The width-height ratio is assumed to be 1. '
)
parser.add_argument(
    '--h_dim', default=64, type=int, help='recurrence hidden state dim.'
)
parser.add_argument(
    '--z_dim', default=64, type=int, help='vae latent var. dim.')
parser.add_argument(
    '--neuron_dim', default=1, type=int, help='readout neuron dim')
parser.add_argument('--fullsnn', default=1, type=int, help='model type')

parser.add_argument('--model', default='tecos', type=str, help='model type')
# FOR SURROGATE GRADIENT FUNC
parser.add_argument('--alpha', default=1.0, type=float, 
                    help='surrogate grad func scale hyperparameter')
# Beta-VAE, the lagrangian multiplier
parser.add_argument('--beta', default=1.0, type=float, 
                    help='beta for beta-vae, 1 for vanilla vae')
# Vanilla regularizers
parser.add_argument('--weight_decay', default=0.0, type=float, 
                    help='lagrangian factor of L2 norm term')

# Other opt. settings
parser.add_argument(
    '--grad_clip_thresh', default=20, type=float, 
    help='gradient clipping threshold for stable training.'
)
parser.add_argument('--temperature', default=1.0, type=float, help='temperature')
parser.add_argument('--seed', default=1000, type=int, help='random seed')
parser.add_argument('--workers', default=8, type=int, help='#threads')
parser.add_argument('--optim', default='adam', type=str, help='optimizer')
parser.add_argument('--scheduler', default='cos', type=str, help='lr scheduler')

# Experimental functions.
parser.add_argument('--init', default='uniform', type=str, 
                    help="specified weight initialization method. ")                
parser.add_argument('--plot_spike_rate', default=False, type=bool, 
                    help="plot spike rate or not. ") 
# 
parser.add_argument('--resume', '-r', action='store_true',
                    help='resume from checkpoint')
parser.add_argument(
    '--debug', action='store_true',
    help='if debug is true, formal log files will not be created'
)

# dataset specification
parser.add_argument(
    '--exp', default=1, type=int, 
    help='which exp. They use two different retinas in exp 1 and exp 2.'
)
parser.add_argument('--mov', default=1, type=int, help='which movie.')
parser.add_argument('--stim_hist', default=0, type=int, help='stimuli history len')
args = parser.parse_args()
print(COLOR_PURPLE+'==> Natural Movie dataset: Exp {}, Movie {}'.format(args.exp, args.mov))
# Set readout neuron dim 
if args.exp == 1:
    args.neuron_dim = 38
elif args.exp == 2:
    args.neuron_dim = 49
##############################################################
basic_neuron = importlib.import_module('models.' + args.neuron).Neuron
    
os.environ['PYTHONHASHSEED'] = str(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True

device = 'cuda' if torch.cuda.is_available() else 'cpu'
best_loss = 9999  
start_epoch = 0  


trainloader, testloader, stim_his, vis_trainloader, vis_testloader = prepare_natural_movie(
    args,
    stim_hist=args.stim_hist
)

print(COLOR_CYAN+'==> Building model..')
if args.model == 'tecos':
    print('==> TeCoS-LVM {} full'.format(args.neuron)+COLOR_END)
    net = TemporalConditioningSpikingLVM(
        x_dim=[args.x_dim, args.x_dim],
        h_dim=args.h_dim,
        z_dim=args.z_dim,
        neuron_dim=args.neuron_dim,
        n_layers=1,
        device=device,
        spiking_neuron=basic_neuron,
        threshold=args.threshold,  # for neuron
        tau=args.tau,  # for neuron 
        alpha=args.alpha,  # for neuron
        sigma=args.sigma,  # for Noisy neuron
    )

net = net.to(device)
if device == 'cuda':
    cudnn.benchmark = True

total_params = sum(param.numel() for param in net.parameters() if param.requires_grad)
print('==> parameter count: {}'.format(total_params))

optimizer = optim.Adam(
    net.parameters(), 
    lr=args.lr, 
    weight_decay=args.weight_decay
)

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer, 
    T_max=args.epochs
) 


def train(epoch, writer, scheduler, args):
    print('\nEpoch: %d' % epoch)
    net.train()
    
    train_loss_total_all = 0
    train_loss_kld_all   = 0
    train_loss_recon_all = 0

    for batch_idx, (stimuli, response, fr) in enumerate(trainloader):
        stimuli, response = stimuli.to(device), response.to(device)
        stimuli = stimuli.transpose(0, 1)  
        response = response.transpose(0, 1)
        fr = fr.transpose(0, 1).to(device)
        
        stimuli = stimuli.unsqueeze(2) 
        timesteps = stimuli.size(0)
        batch_size = stimuli.size(1) 
            
        optimizer.zero_grad()
        
        net.reset()
      
        kld_loss, recon_loss, _, [pred, _] = net(stimuli, response, fr)
        kld_loss /= batch_size
        recon_loss /= batch_size

        loss_total = recon_loss + args.beta * kld_loss
        loss_total.backward()

        nn.utils.clip_grad_norm_(net.parameters(), args.grad_clip_thresh)
        optimizer.step()

        train_loss_total_all += loss_total.item()
        train_loss_kld_all += kld_loss.item()
        train_loss_recon_all += recon_loss.item()
        
        progress_bar(
            batch_idx, 
            len(trainloader), 
            'Loss_total: %.5f | Loss_kld: %.5f | Loss_recon: %.5f' % (
                loss_total.item(), kld_loss.item(), recon_loss.item() 
            )
        )
        
    writer.add_scalar(
        'loss_total / train', train_loss_total_all / (batch_idx+1), epoch
    )
    writer.add_scalar(
        'loss_kld / train', train_loss_kld_all / (batch_idx+1), epoch
    )
    writer.add_scalar(
        'loss_recon / train', train_loss_recon_all / (batch_idx+1), epoch
    )

    net.eval()
    with torch.no_grad():
        t_counter = 0
        t_counter_max = 450  # 15 s

        pred_len = 15  # 1
        if args.model == 'swspikinglvm':
            pred_len = 1

        prediction = torch.zeros(t_counter_max, response.size(-1))
        true_response = torch.zeros(t_counter_max, response.size(-1))
        
        for batch_idx, (stimuli, response, fr) in enumerate(vis_trainloader):
            if t_counter >= t_counter_max:
                break
            if not t_counter % pred_len:
                stimuli, response = stimuli.to(device), response.to(device)
                stimuli = stimuli.transpose(0, 1)  
                response = response.transpose(0, 1)
                fr = fr.transpose(0, 1).to(device)

                stimuli = stimuli.unsqueeze(2) 
                batch_size = stimuli.size(1) 

                true_response[
                    t_counter: t_counter+pred_len
                ] = response.squeeze()[-pred_len:]

                net.reset()

                kld_loss, recon_loss, _, [pred, _] = net(stimuli, response, fr) 
                prediction[
                    t_counter: t_counter+pred_len
                ] = pred.squeeze()[-pred_len:]

            t_counter += 1 

    writer.add_image('response train / true', true_response.T.unsqueeze(0), epoch)
    writer.add_image('response train / prediction', prediction.T.unsqueeze(0), epoch)
        

def test(epoch, writer, path):
    global best_loss
    net.eval()
    test_loss_total_all = 0
    test_loss_kld_all   = 0
    test_loss_recon_all = 0
    test_cc_all         = 0
    with torch.no_grad():
        for batch_idx, (stimuli, response, fr) in enumerate(testloader):
            stimuli, response = stimuli.to(device), response.to(device)
            stimuli = stimuli.transpose(0, 1)  
            response = response.transpose(0, 1)
            fr = fr.transpose(0, 1).to(device)

            stimuli = stimuli.unsqueeze(2)  
            
            net.reset()

            timesteps = stimuli.size(0)
            batch_size = stimuli.size(1)
            
            kld_loss, recon_loss, _, [pred, _] = net(stimuli, response, fr)
            
            kld_loss /= batch_size
            recon_loss /= batch_size

            loss_total = recon_loss + args.beta * kld_loss

            test_loss_total_all += loss_total.item()
            test_loss_kld_all += kld_loss.item()
            test_loss_recon_all += recon_loss.item()

            progress_bar(
                batch_idx, len(testloader), 
                'Loss_total: %.5f | Loss_kld: %.5f | Loss_recon: %.5f' % (
                    loss_total.item(), kld_loss.item(), recon_loss.item()
                )
            )
            
    test_loss_total = test_loss_total_all / (batch_idx+1)
    test_loss_recon_ = test_loss_recon_all / (batch_idx+1)
    
    writer.add_scalar('loss_total / test', test_loss_total, epoch)
    writer.add_scalar('loss_kld / test', test_loss_kld_all / (batch_idx+1), epoch)
    writer.add_scalar(
        'loss_recon / test', test_loss_recon_all / (batch_idx+1), 
        epoch
    )
    
    with torch.no_grad():
        t_counter = 0
        t_counter_max = 450  # 15 s

        pred_len = 15  # 1
        if args.model == 'swspikinglvm':
            pred_len = 1

        prediction = torch.zeros(t_counter_max, response.size(-1))
        true_response = torch.zeros(t_counter_max, response.size(-1))
        
        for batch_idx, (stimuli, response, fr) in enumerate(vis_trainloader):
            if t_counter >= t_counter_max:
                break
            if not t_counter % pred_len:
                stimuli, response = stimuli.to(device), response.to(device)
                stimuli = stimuli.transpose(0, 1)  
                response = response.transpose(0, 1)
                fr = fr.transpose(0, 1).to(device)

                stimuli = stimuli.unsqueeze(2)  

                true_response[
                    t_counter: t_counter+pred_len
                ] = response.squeeze()[-pred_len:]

                net.reset()
                
                kld_loss, recon_loss, _, [pred, _] = net(stimuli, response, fr) 
                prediction[
                    t_counter: t_counter+pred_len
                ] = pred.squeeze()[-pred_len:]

            t_counter += 1 

    writer.add_image('response test / true', true_response.T.unsqueeze(0), epoch)
    writer.add_image('response test / prediction', prediction.T.unsqueeze(0), epoch)

    if test_loss_recon_ < best_loss:
        print('Saving..')
        state = {
            'net': net.state_dict(),
            'loss': test_loss_recon_,
            'epoch': epoch,
        }
        
        torch.save(
            state, 
            os.path.join(path, 'best.pth')
        )
        best_loss = test_loss_recon_


if __name__ == '__main__':
    uid = uuid.uuid4().hex
    path = os.path.join(
        './results/logs/NatMov-exp_{}-mov_{}-type_{}-neuron_{}-optim_{}-sch_{}-lr_{}-bs_{}-vth_{}-tau_{}-alpha{}-sigma_{}-xdim_{}-hdim_{}-zdim_{}-beta_{}-wd_{}'.format(
            int(args.exp),
            int(args.mov),
            str(args.model),
            str(args.neuron),
            str(args.optim),
            str(args.scheduler),
            float(args.lr), 
            int(args.minibatch),
            float(args.threshold),
            float(args.tau),
            float(args.alpha),
            float(args.sigma),
            int(args.x_dim),
            int(args.h_dim),
            int(args.z_dim),
            float(args.beta), 
            float(args.weight_decay)           
        ),
        str(uid)
    )

    if args.debug:
        writer = SummaryWriter('./')
    else:
        if not os.path.isdir(path):
            os.makedirs(path)
        script_path = os.path.join(path, 'scripts')
        if not os.path.isdir(script_path):
            os.makedirs(script_path)

        writer = SummaryWriter(path)

    files = list(glob.iglob(os.path.join('./', '*.sh'))) + \
            list(glob.iglob(os.path.join('./', '*.py'))) + \
            list(glob.iglob(os.path.join('./models', '*.py')))
    
    for file in files:
        if not os.path.isfile(file):
            continue
        shutil.copy2(
            file, os.path.join(
                script_path, file.replace('models/', 'models_')
            )
        )

    for epoch in range(start_epoch, start_epoch+args.epochs):
        train(epoch, writer, scheduler, args)
        test(epoch, writer, path)
        scheduler.step()
    
    print('done')