spikingjelly.activation_based.model.spiking_vgg 源代码

import torch
import torch.nn as nn
from copy import deepcopy
from .. import functional, neuron, layer

try:
    from torchvision.models.utils import load_state_dict_from_url
except ImportError:
    from torchvision._internally_replaced_utils import load_state_dict_from_url

__all__ = [
    'SpikingVGG',
    'spiking_vgg11','spiking_vgg11_bn',
    'spiking_vgg13','spiking_vgg13_bn',
    'spiking_vgg16','spiking_vgg16_bn',
    'spiking_vgg19','spiking_vgg19_bn',
]

model_urls = {
    "vgg11": "https://download.pytorch.org/models/vgg11-8a719046.pth",
    "vgg13": "https://download.pytorch.org/models/vgg13-19584684.pth",
    "vgg16": "https://download.pytorch.org/models/vgg16-397923af.pth",
    "vgg19": "https://download.pytorch.org/models/vgg19-dcbb9e9d.pth",
    "vgg11_bn": "https://download.pytorch.org/models/vgg11_bn-6002323d.pth",
    "vgg13_bn": "https://download.pytorch.org/models/vgg13_bn-abd245e5.pth",
    "vgg16_bn": "https://download.pytorch.org/models/vgg16_bn-6c64b313.pth",
    "vgg19_bn": "https://download.pytorch.org/models/vgg19_bn-c79401a0.pth",
}

# modified by https://github.com/pytorch/vision/blob/main/torchvision/models/vgg.py

[文档]class SpikingVGG(nn.Module): def __init__(self, cfg, batch_norm=False, norm_layer=None, num_classes=1000, init_weights=True, spiking_neuron: callable = None, **kwargs): super(SpikingVGG, self).__init__() self.features = self.make_layers(cfg=cfg, batch_norm=batch_norm, norm_layer=norm_layer, neuron=spiking_neuron, **kwargs) self.avgpool = layer.AdaptiveAvgPool2d((7, 7)) self.classifier = nn.Sequential( layer.Linear(512 * 7 * 7, 4096), spiking_neuron(**deepcopy(kwargs)), layer.Dropout(), layer.Linear(4096, 4096), spiking_neuron(**deepcopy(kwargs)), layer.Dropout(), layer.Linear(4096, num_classes), ) if init_weights: self._initialize_weights()
[文档] def forward(self, x): x = self.features(x) x = self.avgpool(x) if self.avgpool.step_mode == 's': x = torch.flatten(x, 1) elif self.avgpool.step_mode == 'm': x = torch.flatten(x, 2) x = self.classifier(x) return x
def _initialize_weights(self): for m in self.modules(): if isinstance(m, layer.Conv2d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, layer.BatchNorm2d): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0) elif isinstance(m, layer.Linear): nn.init.normal_(m.weight, 0, 0.01) nn.init.constant_(m.bias, 0)
[文档] @staticmethod def make_layers(cfg, batch_norm=False, norm_layer=None, neuron: callable = None, **kwargs): if norm_layer is None: norm_layer = layer.BatchNorm2d layers = [] in_channels = 3 for v in cfg: if v == 'M': layers += [layer.MaxPool2d(kernel_size=2, stride=2)] else: conv2d = layer.Conv2d(in_channels, v, kernel_size=3, padding=1) if batch_norm: layers += [conv2d, norm_layer(v), neuron(**deepcopy(kwargs))] else: layers += [conv2d, neuron(**deepcopy(kwargs))] in_channels = v return nn.Sequential(*layers)
def sequential_forward(sequential, x_seq): assert isinstance(sequential, nn.Sequential) out = x_seq for i in range(len(sequential)): m = sequential[i] if isinstance(m, neuron.BaseNode): out = m(out) else: out = functional.seq_to_ann_forward(out, m) return out cfgs = { 'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'], 'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'], } def _spiking_vgg(arch, cfg, batch_norm, pretrained, progress, norm_layer: callable = None, spiking_neuron: callable = None, **kwargs): if pretrained: kwargs['init_weights'] = False if batch_norm: norm_layer = norm_layer else: norm_layer = None model = SpikingVGG(cfg=cfgs[cfg], batch_norm=batch_norm, norm_layer=norm_layer, spiking_neuron=spiking_neuron, **kwargs) if pretrained: state_dict = load_state_dict_from_url(model_urls[arch], progress=progress) model.load_state_dict(state_dict) return model
[文档]def spiking_vgg11(pretrained=False, progress=True, spiking_neuron: callable = None, **kwargs): """ :param pretrained: If True, the SNN will load parameters from the ANN pre-trained on ImageNet :type pretrained: bool :param progress: If True, displays a progress bar of the download to stderr :type progress: bool :param spiking_neuron: a spiking neuron layer :type spiking_neuron: callable :param kwargs: kwargs for `spiking_neuron` :type kwargs: dict :return: Spiking VGG-11 :rtype: torch.nn.Module A spiking version of VGG-11 model from `"Very Deep Convolutional Networks for Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_ """ return _spiking_vgg('vgg11', 'A', False, pretrained, progress, None, spiking_neuron, **kwargs)
[文档]def spiking_vgg11_bn(pretrained=False, progress=True, norm_layer: callable = None, spiking_neuron: callable = None, **kwargs): """ :param pretrained: If True, the SNN will load parameters from the ANN pre-trained on ImageNet :type pretrained: bool :param progress: If True, displays a progress bar of the download to stderr :type progress: bool :param norm_layer: a batch norm layer :type norm_layer: callable :param spiking_neuron: a spiking neuron layer :type spiking_neuron: callable :param kwargs: kwargs for `spiking_neuron` :type kwargs: dict :return: Spiking VGG-11 with norm layer :rtype: torch.nn.Module A spiking version of VGG-11-BN model from `"Very Deep Convolutional Networks for Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_ """ return _spiking_vgg('vgg11_bn', 'A', True, pretrained, progress, norm_layer, spiking_neuron, **kwargs)
[文档]def spiking_vgg13(pretrained=False, progress=True, spiking_neuron: callable = None, **kwargs): """ :param pretrained: If True, the SNN will load parameters from the ANN pre-trained on ImageNet :type pretrained: bool :param progress: If True, displays a progress bar of the download to stderr :type progress: bool :param spiking_neuron: a spiking neuron layer :type spiking_neuron: callable :param kwargs: kwargs for `spiking_neuron` :type kwargs: dict :return: Spiking VGG-13 :rtype: torch.nn.Module A spiking version of VGG-13 model from `"Very Deep Convolutional Networks for Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_ """ return _spiking_vgg('vgg13', 'B', False, pretrained, progress, None, spiking_neuron, **kwargs)
[文档]def spiking_vgg13_bn(pretrained=False, progress=True, norm_layer: callable = None, spiking_neuron: callable = None, **kwargs): """ :param pretrained: If True, the SNN will load parameters from the ANN pre-trained on ImageNet :type pretrained: bool :param progress: If True, displays a progress bar of the download to stderr :type progress: bool :param norm_layer: a batch norm layer :type norm_layer: callable :param spiking_neuron: a spiking neuron layer :type spiking_neuron: callable :param kwargs: kwargs for `spiking_neuron` :type kwargs: dict :return: Spiking VGG-11 with norm layer :rtype: torch.nn.Module A spiking version of VGG-13-BN model from `"Very Deep Convolutional Networks for Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_ """ return _spiking_vgg('vgg13_bn', 'B', True, pretrained, progress, norm_layer, spiking_neuron, **kwargs)
[文档]def spiking_vgg16(pretrained=False, progress=True, spiking_neuron: callable = None, **kwargs): """ :param pretrained: If True, the SNN will load parameters from the ANN pre-trained on ImageNet :type pretrained: bool :param progress: If True, displays a progress bar of the download to stderr :type progress: bool :param spiking_neuron: a spiking neuron layer :type spiking_neuron: callable :param kwargs: kwargs for `spiking_neuron` :type kwargs: dict :return: Spiking VGG-16 :rtype: torch.nn.Module A spiking version of VGG-16 model from `"Very Deep Convolutional Networks for Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_ """ return _spiking_vgg('vgg16', 'D', False, pretrained, progress, None, spiking_neuron, **kwargs)
[文档]def spiking_vgg16_bn(pretrained=False, progress=True, norm_layer: callable = None, spiking_neuron: callable = None, **kwargs): """ :param pretrained: If True, the SNN will load parameters from the ANN pre-trained on ImageNet :type pretrained: bool :param progress: If True, displays a progress bar of the download to stderr :type progress: bool :param norm_layer: a batch norm layer :type norm_layer: callable :param spiking_neuron: a spiking neuron layer :type spiking_neuron: callable :param kwargs: kwargs for `spiking_neuron` :type kwargs: dict :return: Spiking VGG-16 with norm layer :rtype: torch.nn.Module A spiking version of VGG-16-BN model from `"Very Deep Convolutional Networks for Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_ """ return _spiking_vgg('vgg16_bn', 'D', True, pretrained, progress, norm_layer, spiking_neuron, **kwargs)
[文档]def spiking_vgg19(pretrained=False, progress=True, spiking_neuron: callable = None, **kwargs): """ :param pretrained: If True, the SNN will load parameters from the ANN pre-trained on ImageNet :type pretrained: bool :param progress: If True, displays a progress bar of the download to stderr :type progress: bool :param spiking_neuron: a spiking neuron layer :type spiking_neuron: callable :param kwargs: kwargs for `spiking_neuron` :type kwargs: dict :return: Spiking VGG-19 :rtype: torch.nn.Module A spiking version of VGG-19 model from `"Very Deep Convolutional Networks for Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_ """ return _spiking_vgg('vgg19', 'E', False, pretrained, progress, None, spiking_neuron, **kwargs)
[文档]def spiking_vgg19_bn(pretrained=False, progress=True, norm_layer: callable = None, spiking_neuron: callable = None, **kwargs): """ :param pretrained: If True, the SNN will load parameters from the ANN pre-trained on ImageNet :type pretrained: bool :param progress: If True, displays a progress bar of the download to stderr :type progress: bool :param norm_layer: a batch norm layer :type norm_layer: callable :param spiking_neuron: a spiking neuron layer :type spiking_neuron: callable :param kwargs: kwargs for `spiking_neuron` :type kwargs: dict :return: Spiking VGG-19 with norm layer :rtype: torch.nn.Module A spiking version of VGG-19-BN model from `"Very Deep Convolutional Networks for Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_ """ return _spiking_vgg('vgg19_bn', 'E', True, pretrained, progress, norm_layer, spiking_neuron, **kwargs)