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(, dim=1)
c4=self.conv4(merge4)
up_5=self.up5(c4)
merge5 = torch.cat(,dim=1)
c5=self.conv5(merge5)
c6=self.conv6(c5)
out = nn.Sigmoid()(c6)
return out
总结：
extra_large的Unet网络结构：模型大小：1944756KB，1.85GB，Size：16~4096
large Unet网络1：模型大小：121652KB， 118MB，Size：16~1024
large Unet网络2：：模型大小：486257KB， 474MB，Size：32~2048
large Unet网络：：模型大小：121329KB， 118MB，Size：64~1024
normal网络：：模型大小：30126KB， 29.4MB，Size：64~512
normal 网络2：：模型大小：7306KB， 7.13MB，Size：64~256
normal 网络3：：模型大小：29126KB， 28.4MB，Size：128~512
normal 网络4：：模型大小：120329KB， 117MB，Size：128~1024
normal 网络5：：模型大小：67176KB， 66.1MB，Size：96~768
small 网络：：模型大小：16400KB， 16MB， Size：96~384
small 网络1：：模型大小：65474KB， 63.9MB，Size：192~768
排序：
extra_large的Unet网络结构：模型大小：1944756KB，1.85GB，Size：16~4096
large Unet网络2：：模型大小：486257KB， 474MB，Size：32~2048
large Unet网络：：模型大小：121329KB， 118MB，Size：64~1024
large Unet网络1：模型大小：121652KB， 118MB，Size：16~1024
normal 网络4：：模型大小：120329KB， 117MB，Size：128~1024
normal 网络3：：模型大小：29126KB， 28.4MB，Size：128~512
normal网络：：模型大小：30126KB， 29.4MB，Size：64~512
normal 网络2：：模型大小：7306KB， 7.13MB，Size：64~256
small 网络1：：模型大小：65474KB， 63.9MB，Size：192~768
normal 网络5：：模型大小：67176KB， 66.1MB，Size：96~768
small 网络：：模型大小：16400KB， 16MB， Size：96~384
按照大小排序：
extra_large的Unet网络结构：模型大小：1944756KB，1.85GB，Size：16~4096
large Unet网络2：：模型大小：486257KB， 474MB，Size：32~2048
large Unet网络：：模型大小：121329KB， 118MB，Size：64~1024 ---- large net
large Unet网络1：模型大小：121652KB， 118MB，Size：16~1024
normal 网络4：：模型大小：120329KB， 117MB，Size：128~1024
small 网络1：：模型大小：65474KB， 63.9MB，Size：192~768
normal 网络5：：模型大小：67176KB， 66.1MB，Size：96~768 ---- normal net
normal网络：：模型大小：30126KB， 29.4MB，Size：64~512 ---- small net
normal 网络3：：模型大小：29126KB， 28.4MB，Size：128~512
small 网络：：模型大小：16400KB， 16MB， Size：96~384
normal 网络2：：模型大小：7306KB， 7.13MB，Size：64~256
页:
1
[2]