|
楼主 |
发表于 2019-6-23 22:21:30
|
显示全部楼层
small 网络1::模型大小:65474KB,63.9MB
- import torch.nn as nn
- import torch
- from torch import autograd
- import torchvision.models as models
- # 导入模型结构
- #resnet50 = models.resnet50(pretrained=True)
- # 加载预先下载好的预训练参数到resnet18
- #resnet50.load_state_dict(torch.load('resnet50-5c106cde.pth'))
- class DoubleConv(nn.Module):
- def __init__(self, in_ch, out_ch):
- super(DoubleConv, self).__init__()
- self.conv = nn.Sequential(
- nn.Conv2d(in_ch, out_ch, 3, padding=1),
- nn.BatchNorm2d(out_ch),
- nn.ReLU(inplace=True),
- nn.Conv2d(out_ch, out_ch, 3, padding=1),
- nn.BatchNorm2d(out_ch),
- nn.ReLU(inplace=True)
- )
- def forward(self, input):
- return self.conv(input)
- class Unet(nn.Module):
- def __init__(self,in_ch,out_ch):
- super(Unet, self).__init__()
- # 256x256
- self.conv1 = DoubleConv(in_ch, 192)
- self.pool1 = nn.MaxPool2d(2)
- self.conv2 = DoubleConv(192, 384)
- self.pool2 = nn.MaxPool2d(2)
- self.conv3 = DoubleConv(384, 768)
- self.up4 = nn.ConvTranspose2d(768, 384, 2, stride=2)
- self.conv4 = DoubleConv(768, 384)
- self.up5 = nn.ConvTranspose2d(384, 192, 2, stride=2)
- self.conv5 = DoubleConv(384, 192)
- self.conv6 = nn.Conv2d(192,out_ch, 1)
- def forward(self,x):
- c1=self.conv1(x)
- p1=self.pool1(c1)
- c2=self.conv2(p1)
- p2=self.pool2(c2)
- c3=self.conv3(p2)
- up_4=self.up4(c3)
- merge4 = torch.cat([up_4, c2], dim=1)
- c4=self.conv4(merge4)
- up_5=self.up5(c4)
- merge5 = torch.cat([up_5, c1],dim=1)
- c5=self.conv5(merge5)
- c6=self.conv6(c5)
- out = nn.Sigmoid()(c6)
- return out
复制代码
|
|