import torch as t

class LeNet(t.nn.Module):

    def __init__(self):

        super(LeNet, self).__init__()

        self.features = t.nn.Sequential(

            t.nn.Conv2d(3, 6, 5),

            t.nn.ReLU(),

            t.nn.MaxPool2d(2, 2),

            t.nn.Conv2d(6, 16, 5),

            t.nn.ReLU(),

            t.nn.MaxPool2d(2, 2)

        )

        # 由于调整shape并不是一个class层，

        # 所以在涉及这种操作（非nn.Module操作）需要拆分为多个模型

        self.classifiter = t.nn.Sequential(

            t.nn.Linear(16*5*5, 120),

            t.nn.ReLU(),

            t.nn.Linear(120, 84),

            t.nn.ReLU(),

            t.nn.Linear(84, 10)

        )

    def forward(self, x):

        x = self.features(x)

        x = x.view(-1, 16*5*5)

        x = self.classifiter(x)

        return x

net = LeNet()

from torch import optim

# 通常的step优化过程

optimizer = optim.SGD(params=net.parameters(), lr=1)

optimizer.zero_grad()  # net.zero_grad()

input_ = t.autograd.Variable(t.randn(1, 3, 32, 32))

output = net(input_)

output.backward(output)

optimizer.step()

# # 直接对不同的网络模块制定不同学习率

optimizer = optim.SGD([{'params': net.features.parameters()}, # 默认lr是1e-5

                       {'params': net.classifiter.parameters(), 'lr': 1e-2}], lr=1e-5)

# # 以层为单位，为不同层指定不同的学习率

# ## 提取指定层对象

special_layers = t.nn.ModuleList([net.classifiter[0], net.classifiter[3]])

# ## 获取指定层参数id

special_layers_params = list(map(id, special_layers.parameters()))

print(special_layers_params)

# ## 获取非指定层的参数id

base_params = filter(lambda p: id(p) not in special_layers_params, net.parameters())

optimizer = t.optim.SGD([{'params': base_params},

                         {'params': special_layers.parameters(), 'lr': 0.01}], lr=0.001)

'''调整学习率'''

# 新建optimizer或者修改optimizer.params_groups对应的学习率

# # 新建optimizer更简单也更推荐，optimizer十分轻量级，所以开销很小

# # 但是新的优化器会初始化动量等状态信息，这对于使用动量的优化器（momentum参数的sgd）可能会造成收敛中的震荡

# ## optimizer.param_groups:长度2的list，optimizer.param_groups[0]：长度6的字典

print(optimizer.param_groups[0]['lr'])

old_lr = 0.1

optimizer = optim.SGD([{'params': net.features.parameters()},

                       {'params': net.classifiter.parameters(), 'lr': old_lr*0.1}], lr=1e-5)