ENet分割笔记

https://github.com/cswhshi/segmentation/blob/master/ENet.py

权重1.21m，但是内存占用较大，2080ti上batch-size为8

  

   

    

     

    
    

     

      # -*- coding: utf-8 -*-
     
    
   

    

     

    
    

     

      import torch
     
    
   

    

     

    
    

     

      import torch.nn as nn
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      """
     
    
   

    

     

    
    

     

      初始阶段模块
     
    
   

    

     

    
    

     

      """
     
    
   

    

     

    
    

     

      class InitialBlock(nn.Module):
     
    
   

    

     

    
    

     
 def __init__(self,in_channels,out_channels,bias=False,relu=True):
     
    
   

    

     

    
    

     

       super(InitialBlock, self).__init__()
     
    
   

    

     

    
    

     
 if(relu):
     
    
   

    

     

    
    

     

       activation=nn.ReLU
     
    
   

    

     

    
    

     
 else:
     
    
   

    

     

    
    

     

       activation=nn.PReLU
     
    
   

    

     

    
    

     
 #主分支
     
    
   

    

     

    
    

     

       self.main_branch=nn.Conv2d(in_channels,out_channels-3,kernel_size=3,stride=2,padding=1,bias=bias)
     
    
   

    

     

    
    

     
 #分支
     
    
  
 

文章来源: blog.csdn.net，作者：网奇，版权归原作者所有，如需转载，请联系作者。

原文链接：blog.csdn.net/jacke121/article/details/103896082