spikezoo 0.1.1__py3-none-any.whl → 0.2__py3-none-any.whl

spikezoo/archs/ssir/losses.py

@@ -0,0 +1,21 @@
+import torch
+import torch.nn.functional as F
+from utils import InputPadder
+
+
+def compute_l1_loss(img_list, gt):
+    l1_loss = 0.0
+    for img in img_list:
+        cur_size = img.shape[-2:]
+        gt_resize = F.interpolate(gt, size=cur_size, mode="bilinear", align_corners=False)
+        l1_loss += (img - gt_resize).abs().mean()
+
+    return l1_loss
+
+
+def compute_per_loss_single(img, gt, vgg):
+    img_relu5_1 = vgg((img.repeat([1,3,1,1]) + 1.) / 2.)
+    with torch.no_grad():
+        gt_relu5_1 = vgg((gt.repeat([1,3,1,1]).detach() + 1.) / 2.)
+    percep_loss = F.mse_loss(img_relu5_1, gt_relu5_1)
+    return percep_loss

spikezoo/archs/ssir/main.py

@@ -0,0 +1,326 @@
+import argparse
+import os
+import os.path as osp
+import time
+import numpy as np
+import torch
+import torch.optim
+import torch.backends.cudnn as cudnn
+from tensorboardX import SummaryWriter
+import pprint
+import datetime
+from configs.yml_parser import *
+from datasets.dataset_sreds import *
+from models.networks import *
+from utils import *
+from metrics.psnr import *
+from metrics.ssim import *
+import lpips
+from losses import *
+from models.Vgg19 import *
+from spikingjelly.clock_driven import functional
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--data-root', '-dr', type=str, default='/home/data/rzhao/REDS_dataset/REDS120fps')
+parser.add_argument('--arch', '-a', type=str, default='SSIR')
+parser.add_argument('--batch-size', '-b', type=int, default=8)
+parser.add_argument('--learning-rate', '-lr', type=float, default=4e-4)
+parser.add_argument('--configs', '-cfg', type=str, default='./configs/SSIR.yml')
+parser.add_argument('--epochs', '-ep', type=int, default=100)
+parser.add_argument('--epoch-size', '-es', type=int, default=1000)
+parser.add_argument('--workers', '-j', type=int, default=8)
+parser.add_argument('--pretrained', '-prt', type=str, default=None)
+parser.add_argument('--start-epoch', '-sep', type=int, default=0)
+parser.add_argument('--print-freq', '-pf', type=int, default=200)
+parser.add_argument('--save-dir', '-sd', type=str, default='outputs')
+parser.add_argument('--save-name', '-sn', type=str, default=None)
+parser.add_argument('--vis-path', '-vp', type=str, default='vis')
+parser.add_argument('--vis-name', '-vn', type=str, default='SSIR')
+parser.add_argument('--eval_path', '-evp', type=str, default='eval_vis')
+parser.add_argument('--vis-freq', '-vf', type=int, default=20)
+parser.add_argument('--eval', '-e', action='store_true')
+parser.add_argument('--w_per', '-wper', type=float, default=0.2)
+parser.add_argument('--print_details', '-pd', action='store_true')
+parser.add_argument('--milestones', default=[20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70], metavar='N', nargs='*')
+parser.add_argument('--lr-scale-factor', '-lrsf', type=float, default=0.7)
+parser.add_argument('--eval-interval', '-ei', type=int, default=5)
+parser.add_argument('--save-interval', '-si', type=int, default=5)
+parser.add_argument('--no_imwrite', action='store_true', default=False)
+args = parser.parse_args()
+
+args.milestones = [int(m) for m in args.milestones]
+print('milstones', args.milestones)
+
+cfg_parser = YAMLParser(args.configs)
+cfg = cfg_parser.config
+
+cfg['data']['root'] = args.data_root
+cfg = add_args_to_cfg(cfg, args, ['batch_size', 'arch', 'learning_rate', 'configs', 'epochs', 'epoch_size', 'workers', 'pretrained', 'start_epoch', 
+                        'print_freq', 'save_dir', 'save_name', 'vis_path', 'vis_name', 'eval_path', 'vis_freq', 'w_per'])
+
+n_iter = 0
+
+
+
+def train(cfg, train_loader, model, optimizer, epoch, train_writer):
+    ######################################################################
+    ## Init
+    global n_iter
+    batch_time = AverageMeter()
+    data_time = AverageMeter()
+    losses_name = ['rec_loss', 'per_loss', 'all_loss']
+    losses = AverageMeter(precision=6, i=len(losses_name), names=losses_name)
+    model.train()
+    end = time.time()
+    
+    vgg19 = Vgg19(requires_grad=False).cuda()
+    if torch.cuda.device_count() > 1:
+        vgg19 = nn.DataParallel(vgg19, list(range(torch.cuda.device_count())))
+
+
+    ######################################################################
+    ## Training Loop
+    
+    for ww, data in enumerate(train_loader, 0):
+        
+        if ww >= args.epoch_size:
+            return
+
+        spikes = [spk.cuda() for spk in data['spikes']]
+        images = [img.cuda() for img in data['images']]
+        data_time.update(time.time() - end)
+
+        cur_spks = torch.cat(spikes, dim=1)
+
+        rec_loss = 0.0
+        per_loss = 0.0
+        
+        for jj in range(1, 1+cfg['model']['seq_len']):
+            x = cur_spks[:, jj*20-20 : jj*20+21]
+
+            gt = images[jj]
+
+            out_list = model(x)
+            rec_list = [torch.clip(out, 0, 1) for out in out_list]
+
+            # if jj > 1+2:
+            if jj >= 2:
+                rec_loss += compute_l1_loss(rec_list, gt) / (cfg['model']['seq_len'] - 2)
+                if cfg['train']['w_per'] > 0:
+                    per_loss += cfg['train']['w_per'] * compute_per_loss_single(rec_list[-1], gt, vgg19) / (cfg['model']['seq_len'] - 2)
+                else:
+                    per_loss = torch.tensor([0.0]).cuda()
+                
+        all_loss = rec_loss + per_loss
+        
+        # record loss
+        losses.update([rec_loss.item(), per_loss.item(), all_loss.item()])
+        train_writer.add_scalar('rec_loss', rec_loss.item(), n_iter)
+        train_writer.add_scalar('per_loss', per_loss.item(), n_iter)
+        train_writer.add_scalar('total_loss', all_loss.item(), n_iter)
+
+        ## compute gradient and optimize
+        all_loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+        functional.reset_net(model)
+        
+        batch_time.update(time.time() - end)
+        end = time.time()
+        n_iter += 1
+
+        if n_iter % cfg['train']['vis_freq'] == 0:
+            vis_img(cfg['train']['vis_path'], rec_list[-1], cfg['train']['vis_name'])
+        
+        if ww % cfg['train']['print_freq'] == 0:
+            out_str = 'Epoch: [{:d}] [{:d}/{:d}],  Iter: {:d}  '.format(epoch, ww, len(train_loader), n_iter-1)
+            out_str += 'Time: {},  Data: {}  '.format(batch_time, data_time)
+            out_str += ' '.join(map('{:s} {:.4f} ({:.6f}) '.format, losses.names, losses.val, losses.avg))
+            out_str += 'lr {:.6f}'.format(optimizer.state_dict()['param_groups'][0]['lr'])
+            print(out_str)
+        
+        end = time.time()
+    
+    return
+
+
+def validation(cfg, test_loader, model):
+    global n_iter
+    batch_time = AverageMeter()
+    data_time = AverageMeter()
+    metrics_name = ['PSNR', 'SSIM', 'LPIPS', 'AvgTime']
+    all_metrics = AverageMeter(i=len(metrics_name), precision=4, names=metrics_name)
+
+    model.eval()
+
+    loss_fn_vgg = lpips.LPIPS(net='alex').cuda()
+
+    for ww, data in enumerate(test_loader, 0):
+        st1 = time.time()
+        spikes = torch.cat([spk.cuda() for spk in data['spikes']], dim=1)
+        images = data['images']
+        data_time.update(time.time() - st1)
+
+        seq_metrics = AverageMeter(i=len(metrics_name), precision=4, names=metrics_name)
+
+        seq_len = len(data['spikes']) - 2
+
+        rec = []
+        for jj in range(1, 1+seq_len):
+            x = spikes[:, jj*20-20 : jj*20+21]
+            
+            gt = images[jj].cuda()
+
+            with torch.no_grad():
+                st = time.time()
+                out = model(x)
+                mtime = time.time() - st
+            rec = torch.clip(out, 0, 1)
+
+            cur_rec = torch2numpy255(rec)    
+            cur_gt = torch2numpy255(gt)
+
+            if not args.no_imwrite:
+                cur_vis_path = osp.join(args.eval_path, '{:03d}_{:03d}.png'.format(ww, jj))
+                cv2.imwrite(cur_vis_path, cur_rec.astype(np.uint8))
+
+            cur_psnr = calculate_psnr(cur_rec, cur_gt)
+            cur_ssim = calculate_ssim(cur_rec, cur_gt)
+            with torch.no_grad():
+                cur_lpips = loss_fn_vgg(rec, gt)
+
+            cur_metrics_list = [cur_psnr, cur_ssim, cur_lpips.item(), mtime]
+
+            all_metrics.update(cur_metrics_list)
+            seq_metrics.update(cur_metrics_list)
+        
+        functional.reset_net(model)
+            
+        if args.print_details:
+            ostr = 'Data{:02d}  '.format(ww) + ' '.join(map('{:s} {:.4f} '.format, seq_metrics.names, seq_metrics.avg))
+            print(ostr)
+            print()
+
+    
+    ostr = 'All  ' + ' '.join(map('{:s} {:.4f} '.format, all_metrics.names, all_metrics.avg))
+    print(ostr)
+
+    return
+    
+
+def main():
+    ##########################################################################################################
+    # Set random seeds
+    set_seeds(cfg['seed'])
+
+    # Create save path and logs
+    timestamp1 = datetime.datetime.now().strftime('%m-%d')
+    timestamp2 = datetime.datetime.now().strftime('%H%M%S')
+    if args.save_name == None:
+        save_folder_name = 'b{:d}_{:s}'.format(args.batch_size, timestamp2)
+    else:
+        save_folder_name = 'b{:d}_{:s}_{:s}'.format(args.batch_size, timestamp2, args.save_name)
+
+    save_path = osp.join(args.save_dir, timestamp1, save_folder_name)
+    print('save path: ', save_path)
+    make_dir(save_path)
+    make_dir(args.vis_path)
+    make_dir(args.eval_path)    
+
+    train_writer = SummaryWriter(save_path)
+
+    cfg_str = pprint.pformat(cfg)
+    print('=> configurations: ')
+    print(cfg_str)
+    
+    ##########################################################################################################
+    ## Create model
+    model = eval(args.arch)()
+
+    if args.pretrained:
+        network_data = torch.load(args.pretrained)
+        print('=> using pretrained model {:s}'.format(args.pretrained))
+        model = torch.nn.DataParallel(model).cuda()
+        model = model.cuda()
+        model.load_state_dict(network_data)
+    else:
+        network_data = None
+        print('=> train from scratch')
+        model.init_weights()
+        print('=> model params: {:.6f}M'.format(model.num_parameters()/1e6))
+        model = torch.nn.DataParallel(model).cuda()
+        model = model.cuda()
+
+    cudnn.benchmark = True
+
+    ##########################################################################################################
+    ## Create Optimizer
+    cfgopt = cfg['optimizer']
+    cfgmdl = cfg['model']
+    assert(cfgopt['solver'] in ['Adam', 'SGD'])
+    print('=> settings {:s} solver'.format(cfgopt['solver']))
+    
+    param_groups = [{'params': model.parameters(), 'weight_decay': cfgmdl['flow_weight_decay']}]
+    if cfgopt['solver'] == 'Adam':
+        optimizer = torch.optim.Adam(param_groups, args.learning_rate, betas=(cfgopt['momentum'], cfgopt['beta']))
+    elif cfgopt['solver'] == 'SGD':
+        optimizer = torch.optim.SGD(param_groups, args.learning_rate, momentum=cfgopt['momentum'])
+
+    ##########################################################################################################
+    ## Dataset
+    train_set = sreds_train(cfg)
+    train_loader = torch.utils.data.DataLoader(
+        train_set,
+        drop_last=False,
+        batch_size=cfg['train']['batch_size'],
+        shuffle=True,
+        num_workers=cfg['train']['workers'],
+        # pin_memory=True
+    )
+
+    test_set = sreds_test(cfg)
+    test_loader = torch.utils.data.DataLoader(
+        test_set,
+        drop_last=False,
+        batch_size=1,
+        shuffle=False,
+        num_workers=cfg['train']['workers']
+    )
+
+    ##########################################################################################################
+    ## Train or Evaluate
+    if args.eval:
+        validation(cfg=cfg, test_loader=test_loader, model=model)
+    else:
+        epoch = cfg['train']['start_epoch']
+        while(True):
+            train(
+                cfg=cfg,
+                train_loader=train_loader,
+                model=model,
+                optimizer=optimizer,
+                epoch=epoch,
+                train_writer=train_writer
+            )
+            epoch += 1
+
+            # scheduler can be added here
+            if epoch in args.milestones:
+                for param_group in optimizer.param_groups:
+                    param_group['lr'] = param_group['lr'] * args.lr_scale_factor
+
+            # save model
+            if epoch % args.save_interval == 0:
+                model_save_name = '{:s}_epoch{:03d}.pth'.format(cfg['model']['arch'], epoch)
+                torch.save(model.state_dict(), osp.join(save_path, model_save_name))
+            
+            # if epoch % 5 == 0:
+            if epoch % args.eval_interval == 0:
+                validation(cfg=cfg, test_loader=test_loader, model=model)
+
+            if epoch >= cfg['train']['epochs']:
+                break
+
+if __name__ == '__main__':
+    main()

spikezoo/archs/ssir/metrics/psnr.py

@@ -0,0 +1,22 @@
+import math
+import numpy as np
+
+# --------------------------------------------
+# PSNR
+# --------------------------------------------
+def calculate_psnr(img1, img2, border=0):
+    # img1 and img2 have range [0, 255]
+    #img1 = img1.squeeze()
+    #img2 = img2.squeeze()
+    if not img1.shape == img2.shape:
+        raise ValueError('Input images must have the same dimensions.')
+    h, w = img1.shape[:2]
+    img1 = img1[border:h-border, border:w-border]
+    img2 = img2[border:h-border, border:w-border]
+
+    img1 = img1.astype(np.float64)
+    img2 = img2.astype(np.float64)
+    mse = np.mean((img1 - img2)**2)
+    if mse == 0:
+        return float('inf')
+    return 20 * math.log10(255.0 / math.sqrt(mse))

spikezoo/archs/ssir/metrics/ssim.py

@@ -0,0 +1,54 @@
+import numpy as np
+import cv2
+
+# --------------------------------------------
+# SSIM
+# --------------------------------------------
+def calculate_ssim(img1, img2, border=0):
+    '''calculate SSIM
+    the same outputs as MATLAB's
+    img1, img2: [0, 255]
+    '''
+    #img1 = img1.squeeze()
+    #img2 = img2.squeeze()
+    if not img1.shape == img2.shape:
+        raise ValueError('Input images must have the same dimensions.')
+    h, w = img1.shape[:2]
+    img1 = img1[border:h-border, border:w-border]
+    img2 = img2[border:h-border, border:w-border]
+
+    if img1.ndim == 2:
+        return ssim(img1, img2)
+    elif img1.ndim == 3:
+        if img1.shape[2] == 3:
+            ssims = []
+            for i in range(3):
+                ssims.append(ssim(img1[:,:,i], img2[:,:,i]))
+            return np.array(ssims).mean()
+        elif img1.shape[2] == 1:
+            return ssim(np.squeeze(img1), np.squeeze(img2))
+    else:
+        raise ValueError('Wrong input image dimensions.')
+
+
+def ssim(img1, img2):
+    C1 = (0.01 * 255)**2
+    C2 = (0.03 * 255)**2
+
+    img1 = img1.astype(np.float64)
+    img2 = img2.astype(np.float64)
+    kernel = cv2.getGaussianKernel(11, 1.5)
+    window = np.outer(kernel, kernel.transpose())
+
+    mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]  # valid
+    mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
+    mu1_sq = mu1**2
+    mu2_sq = mu2**2
+    mu1_mu2 = mu1 * mu2
+    sigma1_sq = cv2.filter2D(img1**2, -1, window)[5:-5, 5:-5] - mu1_sq
+    sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq
+    sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2
+
+    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
+                                                            (sigma1_sq + sigma2_sq + C2))
+    return ssim_map.mean()

spikezoo/archs/ssir/models/Vgg19.py

@@ -0,0 +1,42 @@
+import torch
+import torch.nn as nn
+from torchvision import models
+
+
+class Vgg19(torch.nn.Module):
+    def __init__(self, requires_grad=False, rgb_range=1):
+        super(Vgg19, self).__init__()
+        
+        vgg_pretrained_features = models.vgg19(pretrained=True).features
+
+        self.slice1 = torch.nn.Sequential()
+        for x in range(30):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+
+        if not requires_grad:
+            for param in self.slice1.parameters():
+                param.requires_grad = False
+        
+        vgg_mean = (0.485, 0.456, 0.406)
+        vgg_std = (0.229 * rgb_range, 0.224 * rgb_range, 0.225 * rgb_range)
+        self.sub_mean = MeanShift(rgb_range, vgg_mean, vgg_std)
+
+    def forward(self, X):
+        h = self.sub_mean(X)
+        h_relu5_1 = self.slice1(h)
+        return h_relu5_1
+
+class MeanShift(nn.Conv2d):
+    def __init__(self, rgb_range, rgb_mean, rgb_std, sign=-1):
+        super(MeanShift, self).__init__(3, 3, kernel_size=1)
+        std = torch.Tensor(rgb_std)
+        self.weight.data = torch.eye(3).view(3, 3, 1, 1)
+        self.weight.data.div_(std.view(3, 1, 1, 1))
+        self.bias.data = sign * rgb_range * torch.Tensor(rgb_mean)
+        self.bias.data.div_(std)
+        # self.requires_grad = False
+        self.weight.requires_grad = False
+        self.bias.requires_grad = False
+
+if __name__ == '__main__':
+    vgg19 = Vgg19(requires_grad=False)

spikezoo/archs/ssir/models/layers.py

@@ -0,0 +1,110 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from spikingjelly.clock_driven import layer, neuron, surrogate
+
+backend = 'torch'
+# backend = 'cupy'
+
+def get_neuron_code(neuron_type):
+    if neuron_type == 'IF':
+        neuron_code = 'neuron.IFNode(surrogate_function=surrogate.ATan(), detach_reset=True)'
+    elif neuron_type == 'LIF':
+        neuron_code = 'neuron.LIFNode(surrogate_function=surrogate.ATan(), detach_reset=True)'
+    elif neuron_type == 'PLIF':
+        neuron_code = 'neuron.ParametricLIFNode(init_tau=2.0, surrogate_function=surrogate.ATan(), detach_reset=True)'
+
+    return neuron_code
+
+
+class ConvLayerSNN(nn.Module):
+    def __init__(self, in_ch, out_ch, kernel_size=3, stride=1, padding=1, neuron_type='PLIF'):
+        super().__init__()
+        
+        neuron_code = get_neuron_code(neuron_type)
+
+        self.layer = nn.Sequential(
+        nn.Conv2d(in_ch, out_ch, kernel_size=kernel_size, padding=padding, stride=stride, bias=False),
+        nn.BatchNorm2d(out_ch),
+        eval(neuron_code)
+    )
+
+    def forward(self, x):
+        return self.layer(x)
+
+
+class BottleneckBlockSNN(nn.Module):
+    def __init__(self, in_ch, out_ch, expansion=4, kernel_size=3, padding=1, neuron_type='PLIF'):
+        super().__init__()
+
+        neuron_code = get_neuron_code(neuron_type)
+
+        mid_ch = out_ch // expansion
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_ch, mid_ch, kernel_size=1, padding=0, stride=1, bias=False),
+            nn.BatchNorm2d(mid_ch),
+        )
+        self.sn1 = eval(neuron_code)
+
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(mid_ch, mid_ch, kernel_size=kernel_size, padding=padding, stride=1, bias=False),
+            nn.BatchNorm2d(mid_ch),
+        )
+        self.sn2 = eval(neuron_code)
+
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(mid_ch, out_ch, kernel_size=1, padding=0, stride=1, bias=False),
+            nn.BatchNorm2d(out_ch),
+        )
+        self.sn3 = eval(neuron_code)
+
+    def forward(self, x):
+        out = self.sn1(self.conv1(x))
+        out = self.sn2(self.conv2(out))
+        out = self.sn3(self.conv3(out))
+
+        out = out + x
+        return out
+
+
+class DeConvLayerSNN(nn.Module):
+    def __init__(self, in_ch, out_ch, kernel_size=3, padding=1, neuron_type='PLIF'):
+        super().__init__()
+
+        neuron_code = get_neuron_code(neuron_type)
+
+        self.layer = nn.Sequential(
+        nn.ConvTranspose2d(in_ch, out_ch, kernel_size=4, padding=1, stride=2, bias=False),
+        nn.BatchNorm2d(out_ch),
+        eval(neuron_code)
+    )
+
+    def forward(self, x):
+        return self.layer(x)
+
+
+class PredHead2(nn.Module):
+    def __init__(self, in_ch, out_ch, kernel_size=3, padding=1):
+        super().__init__()
+        self.conv1 = nn.Sequential(
+        nn.Conv2d(in_ch, 32, kernel_size=3, padding=4, stride=1, bias=False),
+        nn.BatchNorm2d(32)
+        )
+        self.conv2 = nn.Sequential(
+        nn.Conv2d(32, 32, kernel_size=3, padding=0, stride=1, bias=False),
+        nn.BatchNorm2d(32)
+        )
+        self.conv3 = nn.Sequential(
+        nn.Conv2d(32, 32, kernel_size=3, padding=0, stride=1, bias=False),
+        nn.BatchNorm2d(32)
+        )
+        self.conv4 = nn.Conv2d(32, 1, kernel_size=3, padding=0, stride=1, bias=False)
+
+    def forward(self, x):
+        x1 = self.conv1(x)
+        x2 = self.conv2(x1)
+        x3 = self.conv3(x2)
+        out = self.conv4(x3)
+        return out
+

spikezoo/archs/ssir/models/networks.py

@@ -0,0 +1,61 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .layers import *
+
+class BasicModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+    
+    ####################################################################################
+    ## Tools functions for neural networks
+    def weight_parameters(self):
+        return [param for name, param in self.named_parameters() if 'weight' in name]
+
+    def bias_parameters(self):
+        return [param for name, param in self.named_parameters() if 'bias' in name]
+
+    def num_parameters(self):
+        return sum([p.data.nelement() if p.requires_grad else 0 for p in self.parameters()])
+    
+    def init_weights(self):
+        for layer in self.named_modules():
+            if isinstance(layer, nn.Conv2d):
+                nn.init.kaiming_normal_(layer.weight)
+                if layer.bias is not None:
+                    nn.init.constant_(layer.bias, 0)
+
+            elif isinstance(layer, nn.ConvTranspose2d):
+                nn.init.kaiming_normal_(layer.weight)
+                if layer.bias is not None:
+                    nn.init.constant_(layer.bias, 0)
+
+
+########################################################################
+class SSIR(BasicModel):
+    def __init__(self):
+        super().__init__()
+        base_ch = 128
+
+        self.static_conv = ConvLayerSNN(in_ch=41, out_ch=base_ch, stride=1)
+
+        self.enc1 = ConvLayerSNN(in_ch=base_ch , out_ch=base_ch , stride=2)
+        self.eres1 = BottleneckBlockSNN(in_ch=base_ch, out_ch=base_ch)
+
+        self.dec3 = DeConvLayerSNN(in_ch=base_ch, out_ch=base_ch//2)
+
+        self.pred3 = PredHead2(in_ch=base_ch//2 , out_ch=1)
+
+    def forward(self, x):
+        # x: B x C x H x W
+        x0 = self.static_conv(x)
+
+        x1 = self.eres1(self.enc1(x0))
+
+        x7 = self.dec3(x1)
+        out3 = self.pred3(x7)
+
+        if self.training:
+            return [out3]
+        else:
+            return out3

spikezoo/archs/ssir/requirements.txt

@@ -0,0 +1,8 @@
+numpy
+torch
+torchvision
+spikingjelly==0.0.0.0.14
+opencv-python
+tensorboardX
+lpips
+pyyaml

spikezoo/archs/ssir/shells/eval_SREDS.sh

@@ -0,0 +1,6 @@
+cd ../ &&
+python3 main.py \
+--data-root your_data_root \
+--arch SSIR \
+--pretrained ./ckpt/SSIR_e80.pth \
+--eval

spikezoo/archs/ssir/shells/train_SSIR.sh

@@ -0,0 +1,12 @@
+cd ../ &&
+python3 main.py \
+--data-root your_data_root \
+--arch SSIR \
+--batch-size 8 \
+--learning-rate 4e-4 \
+--configs ./configs/SSIR.yml \
+--epochs 80 \
+--workers 8 \
+--w_per 0.2 \
+--milestones 20 25 30 35 40 45 50 55 65 70 \
+--lr-scale-factor 0.7 -pf 1

spikezoo/archs/ssir/test.py

@@ -0,0 +1,3 @@
+from models.networks import SSIR
+net = SSIR()
+