File size: 1,915 Bytes
3b333f6 a19b5e0 3b333f6 |
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 |
# -*- coding: utf-8 -*-
"""
Created on Sun Jul 4 15:07:27 2021
@author: AlexandreN
"""
from __future__ import print_function, division
import torch
import torch.nn as nn
import torchvision
class SingleTractionHead(nn.Module):
def __init__(self):
super(SingleTractionHead, self).__init__()
self.head_locs = nn.Sequential(nn.Linear(2048, 1024),
nn.ReLU(),
nn.Dropout(p=0.3),
nn.Linear(1024, 4),
nn.Sigmoid()
)
# Head class should output the logits over the classe
self.head_class = nn.Sequential(nn.Linear(2048, 128),
nn.ReLU(),
nn.Dropout(p=0.3),
nn.Linear(128, 1))
def forward(self, features):
features = features.view(features.size()[0], -1)
y_bbox = self.head_locs(features)
y_class = self.head_class(features)
res = (y_bbox, y_class)
return res
def create_model():
# setup the architecture of the model
feature_extractor = torchvision.models.resnet50(pretrained=True)
model_body = nn.Sequential(*list(feature_extractor.children())[:-1])
for param in model_body.parameters():
param.requires_grad = False
# Parameters of newly constructed modules have requires_grad=True by default
# num_ftrs = model_body.fc.in_features
model_head = SingleTractionHead()
model = nn.Sequential(model_body, model_head)
return model
def load_weights(model, path='model.pt', device_='cpu'):
checkpoint = torch.load(path, map_location=torch.device(device_))
model.load_state_dict(checkpoint)
return model
|