当前位置 主页 > 网站技术 > 代码类 >

    PyTorch实现AlexNet示例

    栏目:代码类 时间:2020-01-14 15:07

    PyTorch: https://github.com/shanglianlm0525/PyTorch-Networks

    import torch
    import torch.nn as nn
    import torchvision
    
    class AlexNet(nn.Module):
      def __init__(self,num_classes=1000):
        super(AlexNet,self).__init__()
        self.feature_extraction = nn.Sequential(
          nn.Conv2d(in_channels=3,out_channels=96,kernel_size=11,stride=4,padding=2,bias=False),
          nn.ReLU(inplace=True),
          nn.MaxPool2d(kernel_size=3,stride=2,padding=0),
          nn.Conv2d(in_channels=96,out_channels=192,kernel_size=5,stride=1,padding=2,bias=False),
          nn.ReLU(inplace=True),
          nn.MaxPool2d(kernel_size=3,stride=2,padding=0),
          nn.Conv2d(in_channels=192,out_channels=384,kernel_size=3,stride=1,padding=1,bias=False),
          nn.ReLU(inplace=True),
          nn.Conv2d(in_channels=384,out_channels=256,kernel_size=3,stride=1,padding=1,bias=False),
          nn.ReLU(inplace=True),
          nn.Conv2d(in_channels=256,out_channels=256,kernel_size=3,stride=1,padding=1,bias=False),
          nn.ReLU(inplace=True),
          nn.MaxPool2d(kernel_size=3, stride=2, padding=0),
        )
        self.classifier = nn.Sequential(
          nn.Dropout(p=0.5),
          nn.Linear(in_features=256*6*6,out_features=4096),
          nn.ReLU(inplace=True),
          nn.Dropout(p=0.5),
          nn.Linear(in_features=4096, out_features=4096),
          nn.ReLU(inplace=True),
          nn.Linear(in_features=4096, out_features=num_classes),
        )
      def forward(self,x):
        x = self.feature_extraction(x)
        x = x.view(x.size(0),256*6*6)
        x = self.classifier(x)
        return x
    
    
    if __name__ =='__main__':
      # model = torchvision.models.AlexNet()
      model = AlexNet()
      print(model)
    
      input = torch.randn(8,3,224,224)
      out = model(input)
      print(out.shape)
    
    

    以上这篇PyTorch实现AlexNet示例就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持IIS7站长之家。