slice-discovery-human-eval/model_utils.py at main · njohnson99/slice-discovery-human-eval · GitHub

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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
import sys
import torch
from torchvision.models import resnet18

from coco_utils import get_loader, ImageDataset, DirectoryDataset

class ModelWrapper():

    def __init__(self, model, transform_mode = 'normalize', feature_hook = None, get_id = None):
        self.model = model
        self.transform_mode = transform_mode
        self.feature_hook = feature_hook
        self.get_id = get_id

    def predict(self, im):
        if len(im.size()) == 3:
            im = torch.unsqueeze(im, 0)
        return self.model(im.cuda()).cpu().data.numpy()

    def predict_dataset(self, files, labels):
        model = self.model
        feature_hook = self.feature_hook
        get_id = self.get_id

        dataset = ImageDataset(files, labels, transform_mode = self.transform_mode, get_names = True)
        dataloader = get_loader(dataset)

        out = {}
        for batch in dataloader:
            x = batch[0].cuda()
            y = batch[1].numpy()
            f = batch[2]

            y_hat = model(x).data.cpu().numpy()
            if feature_hook is not None:
                rep = feature_hook.features.data.cpu().numpy()[:, :, 0, 0] #This may be specific to ResNets

            for i, v in enumerate(f):
                if get_id is not None:
                    v = get_id(v)
                tmp = {}
                tmp['pred'] = y_hat[i, :]
                tmp['pred_prb'] = torch.sigmoid(torch.Tensor(y_hat[i, :])).numpy()
                tmp['label'] = y[i, :]
                if feature_hook is not None:
                    tmp['rep'] = rep[i, :]
                out[v] = tmp

        return out


    def predict_directory(self, directory):
        model = self.model
        feature_hook = self.feature_hook
        get_id = self.get_id

        dataset = DirectoryDataset(directory, transform_mode = self.transform_mode)
        dataloader = get_loader(dataset)

        out = {}
        for batch in dataloader:
            x = batch[0].cuda()
            f = batch[1]

            y_hat = model(x).data.cpu().numpy()
            if feature_hook is not None:
                rep = feature_hook.features.data.cpu().numpy()[:, :, 0, 0] #This may be specific to ResNets

            for i, v in enumerate(f):
                if get_id is not None:
                    v = get_id(v)
                tmp = {}
                tmp['pred'] = y_hat[i, :]
                if feature_hook is not None:
                    tmp['rep'] = rep[i, :]
                out[v] = tmp

        return out

class Features:
    def __init__(self, requires_grad = None):
        self.features = None
        self.requires_grad = requires_grad

    def __call__(self, modules, module_in, module_out):
        if self.requires_grad is not None:
            module_out.requires_grad = self.requires_grad
        self.features = module_out

import torch

class LinearModel(torch.nn.Module):
    def __init__(self, W, b):
        super(LinearModel, self).__init__()
        linear = torch.nn.Linear(W.shape[0], W.shape[1], bias = True)
        linear.weight = torch.nn.Parameter(W)
        linear.bias = torch.nn.Parameter(b)
        self.linear = linear

    def forward(self, x):
        out = self.linear(x)
        return out

def get_model(out_features = 1, mode = 'tune', parent = 'pretrained', randomize = False):
    # Load the model
    model = resnet18(pretrained = (parent == 'pretrained'))
    # Change the classification layer
    model.fc = torch.nn.Linear(in_features = 512, out_features = out_features)
    # Load the in the parent model weights
    if parent != 'pretrained':
        model.load_state_dict(torch.load(parent))
    if randomize:
        model.fc = torch.nn.Linear(in_features = 512, out_features = out_features)
    # Setup the trainable parameters
    if mode == 'tune':
        return model, model.parameters()
    elif mode == 'transfer':
        for param in model.parameters():
            param.requires_grad = False
        model.fc.weight.requires_grad = True
        model.fc.bias.requires_grad = True
        return model, model.fc.parameters()
    elif mode == 'eval':
        for param in model.parameters():
            param.requires_grad = False
        model.eval()
        return model
    else:
        print('ResNet.py: Could not determine trainable parameters')
        sys.exit(0)

def get_features(model):
    feature_hook = Features()
    handle = list(model.modules())[66].register_forward_hook(feature_hook)
    return feature_hook

def get_lm(model, label_indices = None):
    if label_indices is not None:
        lm = LinearModel(model.fc.weight[label_indices, :], model.fc.bias[label_indices])
    else:
        lm = LinearModel(model.fc.weight, model.fc.bias)
    return lm

def set_lm(model, lm, label_indices = None):
    with torch.no_grad():
        if label_indices is not None:
            model.fc.weight[label_indices, :] = lm.linear.weight
            model.fc.bias[label_indices] = lm.linear.bias
        else:
            model.fc.weight = lm.linear.weight
            model.fc.bias = lm.linear.bias