test_oneshot.py 13 KB
Newer Older
1
2
3
4
5
6
7
import argparse
import torch
import torch.nn.functional as F
import pytorch_lightning as pl
import pytest
from torchvision import transforms
from torchvision.datasets import MNIST
8
from torch import nn
9
from torch.utils.data import Dataset, RandomSampler
10

11
import nni
12
13
import nni.retiarii.nn.pytorch as nn
from nni.retiarii import strategy, model_wrapper, basic_unit
14
from nni.retiarii.experiment.pytorch import RetiariiExeConfig, RetiariiExperiment
15
16
from nni.retiarii.evaluator.pytorch.lightning import Classification, Regression, DataLoader
from nni.retiarii.nn.pytorch import LayerChoice, InputChoice, ValueChoice
17
from nni.retiarii.oneshot.pytorch import DartsLightningModule
Yuge Zhang's avatar
Yuge Zhang committed
18
from nni.retiarii.strategy import BaseStrategy
19
20
21
from pytorch_lightning import LightningModule, Trainer

from .test_oneshot_utils import RandomDataset
22
23


24
25
26
pytestmark = pytest.mark.skipif(pl.__version__ < '1.0', reason='Incompatible APIs')


27
28
29
30
31
32
33
34
35
36
class DepthwiseSeparableConv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.depthwise = nn.Conv2d(in_ch, in_ch, kernel_size=3, groups=in_ch)
        self.pointwise = nn.Conv2d(in_ch, out_ch, kernel_size=1)

    def forward(self, x):
        return self.pointwise(self.depthwise(x))


37
@model_wrapper
38
39
class SimpleNet(nn.Module):
    def __init__(self, value_choice=True):
40
41
42
43
44
45
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = LayerChoice([
            nn.Conv2d(32, 64, 3, 1),
            DepthwiseSeparableConv(32, 64)
        ])
46
47
48
49
        self.dropout1 = LayerChoice([
            nn.Dropout(.25),
            nn.Dropout(.5),
            nn.Dropout(.75)
50
        ])
51
52
53
54
55
56
57
        self.dropout2 = nn.Dropout(0.5)
        if value_choice:
            hidden = nn.ValueChoice([32, 64, 128])
        else:
            hidden = 64
        self.fc1 = nn.Linear(9216, hidden)
        self.fc2 = nn.Linear(hidden, 10)
58
        self.rpfc = nn.Linear(10, 10)
59
        self.input_ch = InputChoice(2, 1)
60
61
62
63

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(self.conv2(x), 2)
64
65
66
67
68
69
70
        x = torch.flatten(self.dropout1(x), 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        x1 = self.rpfc(x)
        x = self.input_ch([x, x1])
71
72
73
74
        output = F.log_softmax(x, dim=1)
        return output


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
@model_wrapper
class MultiHeadAttentionNet(nn.Module):
    def __init__(self, head_count):
        super().__init__()
        embed_dim = ValueChoice(candidates=[32, 64])
        self.linear1 = nn.Linear(128, embed_dim)
        self.mhatt = nn.MultiheadAttention(embed_dim, head_count)
        self.linear2 = nn.Linear(embed_dim, 1)

    def forward(self, batch):
        query, key, value = batch
        q, k, v = self.linear1(query), self.linear1(key), self.linear1(value)
        output, _ = self.mhatt(q, k, v, need_weights=False)
        y = self.linear2(output)
        return F.relu(y)


@model_wrapper
class ValueChoiceConvNet(nn.Module):
    def __init__(self):
        super().__init__()
        ch1 = ValueChoice([16, 32])
        kernel = ValueChoice([3, 5])
        self.conv1 = nn.Conv2d(1, ch1, kernel, padding=kernel // 2)
        self.batch_norm = nn.BatchNorm2d(ch1)
        self.conv2 = nn.Conv2d(ch1, 64, 3)
        self.dropout1 = LayerChoice([
            nn.Dropout(.25),
            nn.Dropout(.5),
            nn.Dropout(.75)
        ])
        self.fc = nn.Linear(64, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.batch_norm(x)
        x = F.relu(x)
        x = F.max_pool2d(self.conv2(x), 2)
        x = torch.mean(x, (2, 3))
        x = self.fc(x)
        return F.log_softmax(x, dim=1)


@model_wrapper
class RepeatNet(nn.Module):
    def __init__(self):
        super().__init__()
        ch1 = ValueChoice([16, 32])
        kernel = ValueChoice([3, 5])
        self.conv1 = nn.Conv2d(1, ch1, kernel, padding=kernel // 2)
        self.batch_norm = nn.BatchNorm2d(ch1)
        self.conv2 = nn.Conv2d(ch1, 64, 3, padding=1)
        self.dropout1 = LayerChoice([
            nn.Dropout(.25),
            nn.Dropout(.5),
            nn.Dropout(.75)
        ])
        self.fc = nn.Linear(64, 10)
        self.rpfc = nn.Repeat(nn.Linear(10, 10), (1, 4))

    def forward(self, x):
        x = self.conv1(x)
        x = self.batch_norm(x)
        x = F.relu(x)
        x = F.max_pool2d(self.conv2(x), 2)
        x = torch.mean(x, (2, 3))
        x = self.fc(x)
        x = self.rpfc(x)
        return F.log_softmax(x, dim=1)


146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
@model_wrapper
class CellNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.stem = nn.Conv2d(1, 5, 7, stride=4)
        self.cells = nn.Repeat(
            lambda index: nn.Cell({
                'conv1': lambda _, __, inp: nn.Conv2d(
                    (5 if index == 0 else 3 * 4) if inp is not None and inp < 1 else 4, 4, 1
                ),
                'conv2': lambda _, __, inp: nn.Conv2d(
                    (5 if index == 0 else 3 * 4) if inp is not None and inp < 1 else 4, 4, 3, padding=1
                ),
            }, 3, merge_op='loose_end'), (1, 3)
        )
        self.fc = nn.Linear(3 * 4, 10)

    def forward(self, x):
        x = self.stem(x)
        x = self.cells(x)
        x = torch.mean(x, (2, 3))
        x = self.fc(x)
        return F.log_softmax(x, dim=1)


171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
@basic_unit
class MyOp(nn.Module):
    def __init__(self, some_ch):
        super().__init__()
        self.some_ch = some_ch
        self.batch_norm = nn.BatchNorm2d(some_ch)

    def forward(self, x):
        return self.batch_norm(x)


@model_wrapper
class CustomOpValueChoiceNet(nn.Module):
    def __init__(self):
        super().__init__()
        ch1 = ValueChoice([16, 32])
        kernel = ValueChoice([3, 5])
        self.conv1 = nn.Conv2d(1, ch1, kernel, padding=kernel // 2)
        self.batch_norm = MyOp(ch1)
        self.conv2 = nn.Conv2d(ch1, 64, 3, padding=1)
        self.dropout1 = LayerChoice([
            nn.Dropout(.25),
            nn.Dropout(.5),
            nn.Dropout(.75)
        ])
        self.fc = nn.Linear(64, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.batch_norm(x)
        x = F.relu(x)
        x = F.max_pool2d(self.conv2(x), 2)
        x = torch.mean(x, (2, 3))
        x = self.fc(x)
        return F.log_softmax(x, dim=1)


208
def _mnist_net(type_, evaluator_kwargs):
209
210
211
212
213
214
215
216
    if type_ == 'simple':
        base_model = SimpleNet(False)
    elif type_ == 'simple_value_choice':
        base_model = SimpleNet()
    elif type_ == 'value_choice':
        base_model = ValueChoiceConvNet()
    elif type_ == 'repeat':
        base_model = RepeatNet()
217
218
    elif type_ == 'cell':
        base_model = CellNet()
219
220
221
222
    elif type_ == 'custom_op':
        base_model = CustomOpValueChoiceNet()
    else:
        raise ValueError(f'Unsupported type: {type_}')
223

224
    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
Yuge Zhang's avatar
Yuge Zhang committed
225
    train_dataset = nni.trace(MNIST)('data/mnist', download=True, train=True, transform=transform)
226
    # Multi-GPU combined dataloader will break this subset sampler. Expected though.
227
228
    train_random_sampler = nni.trace(RandomSampler)(train_dataset, True, int(len(train_dataset) / 20))
    train_loader = nni.trace(DataLoader)(train_dataset, 64, sampler=train_random_sampler)
Yuge Zhang's avatar
Yuge Zhang committed
229
    valid_dataset = nni.trace(MNIST)('data/mnist', download=True, train=False, transform=transform)
230
231
    valid_random_sampler = nni.trace(RandomSampler)(valid_dataset, True, int(len(valid_dataset) / 20))
    valid_loader = nni.trace(DataLoader)(valid_dataset, 64, sampler=valid_random_sampler)
232
    evaluator = Classification(train_dataloader=train_loader, val_dataloaders=valid_loader, **evaluator_kwargs)
233
234
235
236

    return base_model, evaluator


237
def _multihead_attention_net(evaluator_kwargs):
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
    base_model = MultiHeadAttentionNet(1)

    class AttentionRandDataset(Dataset):
        def __init__(self, data_shape, gt_shape, len) -> None:
            super().__init__()
            self.datashape = data_shape
            self.gtshape = gt_shape
            self.len = len

        def __getitem__(self, index):
            q = torch.rand(self.datashape)
            k = torch.rand(self.datashape)
            v = torch.rand(self.datashape)
            gt = torch.rand(self.gtshape)
            return (q, k, v), gt

        def __len__(self):
            return self.len
256

257
258
259
260
    train_set = AttentionRandDataset((1, 128), (1, 1), 1000)
    val_set = AttentionRandDataset((1, 128), (1, 1), 500)
    train_loader = DataLoader(train_set, batch_size=32)
    val_loader = DataLoader(val_set, batch_size=32)
261

262
    evaluator = Regression(train_dataloader=train_loader, val_dataloaders=val_loader, **evaluator_kwargs)
263
    return base_model, evaluator
264
265


266
267
268
269
270
271
272
273
274
275
276
def _test_strategy(strategy_, support_value_choice=True, multi_gpu=False):
    evaluator_kwargs = {
        'max_epochs': 1
    }
    if multi_gpu:
        evaluator_kwargs.update(
            strategy='ddp',
            accelerator='gpu',
            devices=torch.cuda.device_count()
        )

277
278
    to_test = [
        # (model, evaluator), support_or_net
279
280
281
        (_mnist_net('simple', evaluator_kwargs), True),
        (_mnist_net('simple_value_choice', evaluator_kwargs), support_value_choice),
        (_mnist_net('value_choice', evaluator_kwargs), support_value_choice),
282
        (_mnist_net('repeat', evaluator_kwargs), support_value_choice),      # no strategy supports repeat currently
283
284
        (_mnist_net('custom_op', evaluator_kwargs), False),   # this is definitely a NO
        (_multihead_attention_net(evaluator_kwargs), support_value_choice),
285
    ]
286

287
    for (base_model, evaluator), support_or_not in to_test:
Yuge Zhang's avatar
Yuge Zhang committed
288
289
290
291
292
293
        if isinstance(strategy_, BaseStrategy):
            strategy = strategy_
        else:
            strategy = strategy_(base_model, evaluator)
        print('Testing:', type(strategy).__name__, type(base_model).__name__, type(evaluator).__name__, support_or_not)
        experiment = RetiariiExperiment(base_model, evaluator, strategy=strategy)
294

295
296
        config = RetiariiExeConfig()
        config.execution_engine = 'oneshot'
297

298
299
300
301
302
303
        if support_or_not:
            experiment.run(config)
            assert isinstance(experiment.export_top_models()[0], dict)
        else:
            with pytest.raises(TypeError, match='not supported'):
                experiment.run(config)
304
305


306
def test_darts():
307
    _test_strategy(strategy.DARTS())
308
309


310
311
312
313
314
@pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.device_count() <= 1, reason='Must have multiple GPUs.')
def test_darts_multi_gpu():
    _test_strategy(strategy.DARTS(), multi_gpu=True)


315
def test_proxyless():
316
    _test_strategy(strategy.Proxyless(), False)
317
318
319


def test_enas():
Yuge Zhang's avatar
Yuge Zhang committed
320
321
322
323
324
325
    def strategy_fn(base_model, evaluator):
        if isinstance(base_model, MultiHeadAttentionNet):
            return strategy.ENAS(reward_metric_name='val_mse')
        return strategy.ENAS(reward_metric_name='val_acc')

    _test_strategy(strategy_fn)
326
327


328
329
330
331
332
333
334
335
336
337
@pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.device_count() <= 1, reason='Must have multiple GPUs.')
def test_enas_multi_gpu():
    def strategy_fn(base_model, evaluator):
        if isinstance(base_model, MultiHeadAttentionNet):
            return strategy.ENAS(reward_metric_name='val_mse')
        return strategy.ENAS(reward_metric_name='val_acc')

    _test_strategy(strategy_fn, multi_gpu=True)


338
def test_random():
339
    _test_strategy(strategy.RandomOneShot())
340
341


342
343
def test_gumbel_darts():
    _test_strategy(strategy.GumbelDARTS())
344
345


346
347
def test_optimizer_lr_scheduler():
    learning_rates = []
348

349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
    class CustomLightningModule(LightningModule):
        def __init__(self):
            super().__init__()
            self.layer1 = nn.Linear(32, 2)
            self.layer2 = nn.LayerChoice([nn.Linear(2, 2), nn.Linear(2, 2, bias=False)])

        def forward(self, x):
            return self.layer2(self.layer1(x))

        def configure_optimizers(self):
            opt1 = torch.optim.SGD(self.layer1.parameters(), lr=0.1)
            opt2 = torch.optim.Adam(self.layer2.parameters(), lr=0.2)
            return [opt1, opt2], [torch.optim.lr_scheduler.StepLR(opt1, step_size=2, gamma=0.1)]

        def training_step(self, batch, batch_idx):
            loss = self(batch).sum()
            self.log('train_loss', loss)
            return {'loss': loss}

        def on_train_epoch_start(self) -> None:
            learning_rates.append(self.optimizers()[0].param_groups[0]['lr'])

        def validation_step(self, batch, batch_idx):
            loss = self(batch).sum()
            self.log('valid_loss', loss)

        def test_step(self, batch, batch_idx):
            loss = self(batch).sum()
            self.log('test_loss', loss)

    train_data = RandomDataset(32, 32)
    valid_data = RandomDataset(32, 16)

    model = CustomLightningModule()
    darts_module = DartsLightningModule(model, gradient_clip_val=5)
    trainer = Trainer(max_epochs=10)
    trainer.fit(
        darts_module,
        dict(train=DataLoader(train_data, batch_size=8), val=DataLoader(valid_data, batch_size=8))
    )

    assert len(learning_rates) == 10 and abs(learning_rates[0] - 0.1) < 1e-5 and \
        abs(learning_rates[2] - 0.01) < 1e-5 and abs(learning_rates[-1] - 1e-5) < 1e-6