Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import torch
- import torch.nn as nn
- import torch.nn.functional as F
- from collections import OrderedDict
- import numpy as np
- import os
- class Flatten(nn.Module):
- def __init__(self):
- super(Flatten, self).__init__()
- def forward(self, x):
- """
- Arguments:
- x: a float tensor with shape [batch_size, c, h, w].
- Returns:
- a float tensor with shape [batch_size, c*h*w].
- """
- # without this pretrained model isn't working
- x = x.transpose(3, 2).contiguous()
- return x.view(x.size(0), -1)
- class PNet(nn.Module):
- def __init__(self):
- super(PNet, self).__init__()
- # suppose we have input with size HxW, then
- # after first layer: H - 2,
- # after pool: ceil((H - 2)/2),
- # after second conv: ceil((H - 2)/2) - 2,
- # after last conv: ceil((H - 2)/2) - 4,
- # and the same for W
- self.features = nn.Sequential(OrderedDict([
- ('conv1', nn.Conv2d(3, 10, 3, 1)),
- ('prelu1', nn.PReLU(10)),
- ('pool1', nn.MaxPool2d(2, 2, ceil_mode=True)),
- ('conv2', nn.Conv2d(10, 16, 3, 1)),
- ('prelu2', nn.PReLU(16)),
- ('conv3', nn.Conv2d(16, 32, 3, 1)),
- ('prelu3', nn.PReLU(32))
- ]))
- self.conv4_1 = nn.Conv2d(32, 2, 1, 1)
- self.conv4_2 = nn.Conv2d(32, 4, 1, 1)
- pnetPath = os.path.join( os.path.dirname(__file__), 'weights/pnet.npy')
- weights = np.load(pnetPath, allow_pickle=True)[()]
- for n, p in self.named_parameters():
- p.data = torch.FloatTensor(weights[n])
- def forward(self, x):
- """
- Arguments:
- x: a float tensor with shape [batch_size, 3, h, w].
- Returns:
- b: a float tensor with shape [batch_size, 4, h', w'].
- a: a float tensor with shape [batch_size, 2, h', w'].
- """
- x = self.features(x)
- a = self.conv4_1(x)
- b = self.conv4_2(x)
- a = F.softmax(a)
- return b, a
- class RNet(nn.Module):
- def __init__(self):
- super(RNet, self).__init__()
- self.features = nn.Sequential(OrderedDict([
- ('conv1', nn.Conv2d(3, 28, 3, 1)),
- ('prelu1', nn.PReLU(28)),
- ('pool1', nn.MaxPool2d(3, 2, ceil_mode=True)),
- ('conv2', nn.Conv2d(28, 48, 3, 1)),
- ('prelu2', nn.PReLU(48)),
- ('pool2', nn.MaxPool2d(3, 2, ceil_mode=True)),
- ('conv3', nn.Conv2d(48, 64, 2, 1)),
- ('prelu3', nn.PReLU(64)),
- ('flatten', Flatten()),
- ('conv4', nn.Linear(576, 128)),
- ('prelu4', nn.PReLU(128))
- ]))
- self.conv5_1 = nn.Linear(128, 2)
- self.conv5_2 = nn.Linear(128, 4)
- rnetPath = os.path.join( os.path.dirname(__file__), 'weights/rnet.npy')
- weights = np.load(rnetPath)[()]
- for n, p in self.named_parameters():
- p.data = torch.FloatTensor(weights[n])
- def forward(self, x):
- """
- Arguments:
- x: a float tensor with shape [batch_size, 3, h, w].
- Returns:
- b: a float tensor with shape [batch_size, 4].
- a: a float tensor with shape [batch_size, 2].
- """
- x = self.features(x)
- a = self.conv5_1(x)
- b = self.conv5_2(x)
- a = F.softmax(a)
- return b, a
- class ONet(nn.Module):
- def __init__(self):
- super(ONet, self).__init__()
- self.features = nn.Sequential(OrderedDict([
- ('conv1', nn.Conv2d(3, 32, 3, 1)),
- ('prelu1', nn.PReLU(32)),
- ('pool1', nn.MaxPool2d(3, 2, ceil_mode=True)),
- ('conv2', nn.Conv2d(32, 64, 3, 1)),
- ('prelu2', nn.PReLU(64)),
- ('pool2', nn.MaxPool2d(3, 2, ceil_mode=True)),
- ('conv3', nn.Conv2d(64, 64, 3, 1)),
- ('prelu3', nn.PReLU(64)),
- ('pool3', nn.MaxPool2d(2, 2, ceil_mode=True)),
- ('conv4', nn.Conv2d(64, 128, 2, 1)),
- ('prelu4', nn.PReLU(128)),
- ('flatten', Flatten()),
- ('conv5', nn.Linear(1152, 256)),
- ('drop5', nn.Dropout(0.25)),
- ('prelu5', nn.PReLU(256)),
- ]))
- self.conv6_1 = nn.Linear(256, 2)
- self.conv6_2 = nn.Linear(256, 4)
- self.conv6_3 = nn.Linear(256, 10)
- onetPath = os.path.join( os.path.dirname(__file__), 'weights/onet.npy')
- weights = np.load( onetPath )[()]
- for n, p in self.named_parameters():
- p.data = torch.FloatTensor(weights[n])
- def forward(self, x):
- """
- Arguments:
- x: a float tensor with shape [batch_size, 3, h, w].
- Returns:
- c: a float tensor with shape [batch_size, 10].
- b: a float tensor with shape [batch_size, 4].
- a: a float tensor with shape [batch_size, 2].
- """
- x = self.features(x)
- a = self.conv6_1(x)
- b = self.conv6_2(x)
- c = self.conv6_3(x)
- a = F.softmax(a)
- return c, b, a
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement