Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_dilated_resnet.py 2.02 KB
Newer Older
Hang Zhang's avatar
v0.1.0  
Hang Zhang committed
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 
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
## Created by: Hang Zhang
## ECE Department, Rutgers University
## Email: zhang.hang@rutgers.edu
## Copyright (c) 2017
##
## This source code is licensed under the MIT-style license found in the
## LICENSE file in the root directory of this source tree 
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

import os
import numpy as np
import torch
import torch.nn as nn
from torch.autograd import Variable

import encoding.dilated as dresnet
import torchvision.models as orgresnet

class Dilated_ResNet(nn.Module):
    def __init__(self, nclass):
        super(Dilated_ResNet, self).__init__()
        self.pretrained = dresnet.resnet50(pretrained=True)

    def forward(self, x):
        # pre-trained ResNet feature
        x = self.pretrained.conv1(x)
        x = self.pretrained.bn1(x)
        x = self.pretrained.relu(x)
        x = self.pretrained.maxpool(x)

        x = self.pretrained.layer1(x)
        x = self.pretrained.layer2(x)
        x = self.pretrained.layer3(x)
        x = self.pretrained.layer4(x)
        return x
        

class Org_ResNet(nn.Module):
    def __init__(self, nclass):
        super(Org_ResNet, self).__init__()
        self.pretrained = orgresnet.resnet50(pretrained=True)

    def forward(self, x):
        # pre-trained ResNet feature
        x = self.pretrained.conv1(x)
        x = self.pretrained.bn1(x)
        x = self.pretrained.relu(x)
        x = self.pretrained.maxpool(x)

        x = self.pretrained.layer1(x)
        x = self.pretrained.layer2(x)
        x = self.pretrained.layer3(x)
        x = self.pretrained.layer4(x)
        return x
        
def test_resnet():   
    # test the model
    model1 = Dilated_ResNet(10).eval().cuda()
    model2 = Org_ResNet(10).eval().cuda()
    model1.eval()
    model2.eval()
    x = Variable(torch.Tensor(1,3, 224, 224).uniform_(-0.5,0.5)).cuda()
    y1 = model1(x)
    y2 = model2(x)
    print(y1[0][1])
    print(y2[0][1])

if __name__ == "__main__":
    test_resnet()