|
- # Author: Acer Zhang
- # Datetime: 2021/3/19
- # Copyright belongs to the author.
- # Please indicate the source for reprinting.
-
- # Author: Acer Zhang
- # Datetime: 2021/2/25
- # Copyright belongs to the author.
- # Please indicate the source for reprinting.
-
- import paddle
- import paddle.nn as nn
- from paddle.vision.transforms import Compose, ToTensor
- from paddle.vision.datasets import MNIST
-
- # 导入RIFLE模块
- from paddle_rifle.rifle import RIFLECallback
-
-
- class Net(nn.Layer):
- def __init__(self, num_classes=10):
- super(Net, self).__init__()
- self.conv1 = nn.Conv2D(1, 3, 3)
- self.mp = nn.MaxPool2D(2)
- self.conv2 = nn.Conv2D(3, 16, 3)
- self.mp2 = nn.MaxPool2D(2)
- self.fc1 = nn.Linear(400, 100)
- self.fc2 = nn.Linear(100, num_classes)
-
- def forward(self, inputs):
- x = self.conv1(inputs)
- x = self.mp(x)
- x = self.conv2(x)
- x = self.mp2(x)
- x = paddle.flatten(x, 1)
- x = self.fc1(x)
- x = self.fc2(x)
- return x
-
-
- def main(use_init: bool = False):
- transform = Compose([ToTensor()])
-
- train_data = MNIST(transform=transform)
- test_data = MNIST(mode="test", transform=transform)
-
- net = Net(num_classes=10)
- fc_layer = net.fc2
-
- model = paddle.Model(network=net,
- inputs=paddle.static.InputSpec([1, 28, 28], name="ipt"),
- labels=paddle.static.InputSpec([1], dtype="int64", name="lab"))
-
- rifle_cb = RIFLECallback(fc_layer,
- re_init_epoch=1,
- max_re_num=3,
- weight_initializer=paddle.nn.initializer.XavierNormal() if use_init else None)
-
- sgd = paddle.optimizer.SGD(parameters=model.parameters())
- loss = paddle.nn.loss.CrossEntropyLoss()
- acc = paddle.metric.Accuracy((1, 5))
- model.prepare(sgd, loss, acc)
-
- # 开始训练并传入RIFLE Callback
- model.fit(train_data,
- test_data,
- batch_size=256,
- epochs=2,
- log_freq=10,
- callbacks=[rifle_cb])
-
-
- if __name__ == '__main__':
- main()
|