pytorch model_vision显示网络结构

最新推荐文章于 2024-01-25 01:51:51 发布

weixin_44576543

最新推荐文章于 2024-01-25 01:51:51 发布

阅读量259

点赞数

分类专栏： pytorch 文章标签： python 深度学习 easydl

本文链接：https://blog.youkuaiyun.com/weixin_44576543/article/details/121415560

版权

1 show the model structure显示定义网络结构

import torch.nn as nn
import torch
from torchsummary import summary
# from models.resnet import mobilenet

def conv_bn(inp, oup, stride = 1):
    return nn.Sequential(
        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
        nn.BatchNorm2d(oup),
        nn.ReLU()

        # nn.ReLU6()
    )
    
def conv_dw(inp, oup, stride = 1):
    return nn.Sequential(
        nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
        nn.BatchNorm2d(inp),
        nn.ReLU(),
        # nn.ReLU6(),

        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
        nn.BatchNorm2d(oup),
        nn.ReLU(),
        # nn.ReLU6(),LeakyReLU
    )

class MobileNetV1(nn.Module):
    def __init__(self):
        super(MobileNetV1, self).__init__()
        self.stage1 = nn.Sequential(
            # 160,160,3 -> 80,80,32
            conv_bn(3, 32, 2), 
            # 80,80,32 -> 80,80,64