Spaces:
Runtime error
Runtime error
File size: 3,557 Bytes
c7f097c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 |
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models.resnet as resnet
import torchvision.models.vgg as vgg
class MultiConv(nn.Module):
def __init__(self, filter_channels):
super(MultiConv, self).__init__()
self.filters = []
for l in range(0, len(filter_channels) - 1):
self.filters.append(
nn.Conv2d(filter_channels[l], filter_channels[l + 1], kernel_size=4, stride=2))
self.add_module("conv%d" % l, self.filters[l])
def forward(self, image):
'''
:param image: [BxC_inxHxW] tensor of input image
:return: list of [BxC_outxHxW] tensors of output features
'''
y = image
# y = F.relu(self.bn0(self.conv0(y)), True)
feat_pyramid = [y]
for i, f in enumerate(self.filters):
y = f(y)
if i != len(self.filters) - 1:
y = F.leaky_relu(y)
# y = F.max_pool2d(y, kernel_size=2, stride=2)
feat_pyramid.append(y)
return feat_pyramid
class Vgg16(torch.nn.Module):
def __init__(self):
super(Vgg16, self).__init__()
vgg_pretrained_features = vgg.vgg16(pretrained=True).features
self.slice1 = torch.nn.Sequential()
self.slice2 = torch.nn.Sequential()
self.slice3 = torch.nn.Sequential()
self.slice4 = torch.nn.Sequential()
self.slice5 = torch.nn.Sequential()
for x in range(4):
self.slice1.add_module(str(x), vgg_pretrained_features[x])
for x in range(4, 9):
self.slice2.add_module(str(x), vgg_pretrained_features[x])
for x in range(9, 16):
self.slice3.add_module(str(x), vgg_pretrained_features[x])
for x in range(16, 23):
self.slice4.add_module(str(x), vgg_pretrained_features[x])
for x in range(23, 30):
self.slice5.add_module(str(x), vgg_pretrained_features[x])
def forward(self, X):
h = self.slice1(X)
h_relu1_2 = h
h = self.slice2(h)
h_relu2_2 = h
h = self.slice3(h)
h_relu3_3 = h
h = self.slice4(h)
h_relu4_3 = h
h = self.slice5(h)
h_relu5_3 = h
return [h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3]
class ResNet(nn.Module):
def __init__(self, model='resnet18'):
super(ResNet, self).__init__()
if model == 'resnet18':
net = resnet.resnet18(pretrained=True)
elif model == 'resnet34':
net = resnet.resnet34(pretrained=True)
elif model == 'resnet50':
net = resnet.resnet50(pretrained=True)
else:
raise NameError('Unknown Fan Filter setting!')
self.conv1 = net.conv1
self.pool = net.maxpool
self.layer0 = nn.Sequential(net.conv1, net.bn1, net.relu)
self.layer1 = net.layer1
self.layer2 = net.layer2
self.layer3 = net.layer3
self.layer4 = net.layer4
def forward(self, image):
'''
:param image: [BxC_inxHxW] tensor of input image
:return: list of [BxC_outxHxW] tensors of output features
'''
y = image
feat_pyramid = []
y = self.layer0(y)
feat_pyramid.append(y)
y = self.layer1(self.pool(y))
feat_pyramid.append(y)
y = self.layer2(y)
feat_pyramid.append(y)
y = self.layer3(y)
feat_pyramid.append(y)
y = self.layer4(y)
feat_pyramid.append(y)
return feat_pyramid
|