-
Notifications
You must be signed in to change notification settings - Fork 42
/
model.py
117 lines (95 loc) · 4.56 KB
/
model.py
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
import torch.nn as nn
from collections import OrderedDict
from torchmeta.modules import (MetaModule, MetaConv2d, MetaBatchNorm2d,
MetaSequential, MetaLinear)
def conv_block(in_channels, out_channels, **kwargs):
return MetaSequential(OrderedDict([
('conv', MetaConv2d(in_channels, out_channels, **kwargs)),
('norm', nn.BatchNorm2d(out_channels, momentum=1.,
track_running_stats=False)),
('relu', nn.ReLU()),
('pool', nn.MaxPool2d(2))
]))
class MetaConvModel(MetaModule):
"""4-layer Convolutional Neural Network architecture from [1].
Parameters
----------
in_channels : int
Number of channels for the input images.
out_features : int
Number of classes (output of the model).
hidden_size : int (default: 64)
Number of channels in the intermediate representations.
feature_size : int (default: 64)
Number of features returned by the convolutional head.
References
----------
.. [1] Finn C., Abbeel P., and Levine, S. (2017). Model-Agnostic Meta-Learning
for Fast Adaptation of Deep Networks. International Conference on
Machine Learning (ICML) (https://arxiv.org/abs/1703.03400)
"""
def __init__(self, in_channels, out_features, hidden_size=64, feature_size=64):
super(MetaConvModel, self).__init__()
self.in_channels = in_channels
self.out_features = out_features
self.hidden_size = hidden_size
self.feature_size = feature_size
self.features = MetaSequential(OrderedDict([
('layer1', conv_block(in_channels, hidden_size, kernel_size=3,
stride=1, padding=1, bias=True)),
('layer2', conv_block(hidden_size, hidden_size, kernel_size=3,
stride=1, padding=1, bias=True)),
('layer3', conv_block(hidden_size, hidden_size, kernel_size=3,
stride=1, padding=1, bias=True)),
('layer4', conv_block(hidden_size, hidden_size, kernel_size=3,
stride=1, padding=1, bias=True))
]))
self.classifier = MetaLinear(feature_size, out_features, bias=True)
def forward(self, inputs, params=None):
features = self.features(inputs, params=self.get_subdict(params, 'features'))
features = features.view((features.size(0), -1))
logits = self.classifier(features, params=self.get_subdict(params, 'classifier'))
return logits
class MetaMLPModel(MetaModule):
"""Multi-layer Perceptron architecture from [1].
Parameters
----------
in_features : int
Number of input features.
out_features : int
Number of classes (output of the model).
hidden_sizes : list of int
Size of the intermediate representations. The length of this list
corresponds to the number of hidden layers.
References
----------
.. [1] Finn C., Abbeel P., and Levine, S. (2017). Model-Agnostic Meta-Learning
for Fast Adaptation of Deep Networks. International Conference on
Machine Learning (ICML) (https://arxiv.org/abs/1703.03400)
"""
def __init__(self, in_features, out_features, hidden_sizes):
super(MetaMLPModel, self).__init__()
self.in_features = in_features
self.out_features = out_features
self.hidden_sizes = hidden_sizes
layer_sizes = [in_features] + hidden_sizes
self.features = MetaSequential(OrderedDict([('layer{0}'.format(i + 1),
MetaSequential(OrderedDict([
('linear', MetaLinear(hidden_size, layer_sizes[i + 1], bias=True)),
('relu', nn.ReLU())
]))) for (i, hidden_size) in enumerate(layer_sizes[:-1])]))
self.classifier = MetaLinear(hidden_sizes[-1], out_features, bias=True)
def forward(self, inputs, params=None):
features = self.features(inputs, params=self.get_subdict(params, 'features'))
logits = self.classifier(features, params=self.get_subdict(params, 'classifier'))
return logits
def ModelConvOmniglot(out_features, hidden_size=64):
return MetaConvModel(1, out_features, hidden_size=hidden_size,
feature_size=hidden_size)
def ModelConvMiniImagenet(out_features, hidden_size=64):
return MetaConvModel(3, out_features, hidden_size=hidden_size,
feature_size=5 * 5 * hidden_size)
def ModelMLPSinusoid(hidden_sizes=[40, 40]):
return MetaMLPModel(1, 1, hidden_sizes)
if __name__ == '__main__':
model = ModelMLPSinusoid()