/
model.py
34 lines (29 loc) · 1.34 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
import torch.nn as nn
from torchmeta.modules import (MetaModule, MetaSequential, MetaConv2d,
MetaBatchNorm2d, MetaLinear)
from torchmeta.modules.utils import get_subdict
def conv3x3(in_channels, out_channels, **kwargs):
return MetaSequential(
MetaConv2d(in_channels, out_channels, kernel_size=3, padding=1, **kwargs),
MetaBatchNorm2d(out_channels, momentum=1., track_running_stats=False),
nn.ReLU(),
nn.MaxPool2d(2)
)
class ConvolutionalNeuralNetwork(MetaModule):
def __init__(self, in_channels, out_features, hidden_size=64):
super(ConvolutionalNeuralNetwork, self).__init__()
self.in_channels = in_channels
self.out_features = out_features
self.hidden_size = hidden_size
self.features = MetaSequential(
conv3x3(in_channels, hidden_size),
conv3x3(hidden_size, hidden_size),
conv3x3(hidden_size, hidden_size),
conv3x3(hidden_size, hidden_size)
)
self.classifier = MetaLinear(hidden_size, out_features)
def forward(self, inputs, params=None):
features = self.features(inputs, params=get_subdict(params, 'features'))
features = features.view((features.size(0), -1))
logits = self.classifier(features, params=get_subdict(params, 'classifier'))
return logits