本文介绍了如何基于pytorch_lightning测试resnet18不同激活方式在CIFAR10数据集上的精度
特别说明:
1.NoActive:没有任何激活函数
2.SparseActivation:只保留topk的激活,其余清零,topk通过训练得到[初衷是想让激活变得稀疏]
3.SelectiveActive:通过训练得到使用的激活函数
可参考的代码片段：
1.pytorch_lightning 如何使用
2.pytorch如何替换激活函数
3.如何对自定义权值做衰减

一.曲线

1.train_acc

2.val_acc

3.train_loss

4.lr

二.代码

from __future__ import print_function
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts
import pytorch_lightning as pl
from torch.utils.data import DataLoader
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import torchvision.models as models
import os
import numpy as np
from pytorch_lightning.loggers import TensorBoardLogger#torch.set_float32_matmul_precision('medium')class ResidualBlock(nn.Module):def __init__(self, inchannel, outchannel, stride=1):super(ResidualBlock, self).__init__()self.left = nn.Sequential(nn.Conv2d(inchannel, outchannel, kernel_size=3, stride=stride, padding=1, bias=False),nn.BatchNorm2d(outchannel),nn.ReLU(inplace=True),nn.Conv2d(outchannel, outchannel, kernel_size=3, stride=1, padding=1, bias=False),nn.BatchNorm2d(outchannel))self.shortcut = nn.Sequential()if stride != 1 or inchannel != outchannel:self.shortcut = nn.Sequential(nn.Conv2d(inchannel, outchannel, kernel_size=1, stride=stride, bias=False),nn.BatchNorm2d(outchannel))self.act=nn.ReLU()def forward(self, x):out = self.left(x)out += self.shortcut(x)out = self.act(out)return outclass ResNet(nn.Module):def __init__(self, ResidualBlock, num_classes=10):super(ResNet, self).__init__()self.inchannel = 64self.conv1 = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False),nn.BatchNorm2d(64),nn.ReLU(),)self.layer1 = self.make_layer(ResidualBlock, 64,  2, stride=1)self.layer2 = self.make_layer(ResidualBlock, 128, 2, stride=2)self.layer3 = self.make_layer(ResidualBlock, 256, 2, stride=2)self.layer4 = self.make_layer(ResidualBlock, 512, 2, stride=2)self.fc = nn.Linear(512, num_classes)self.dropout=nn.Dropout(0.5)def make_layer(self, block, channels, num_blocks, stride):strides = [stride] + [1] * (num_blocks - 1)layers = []for stride in strides:layers.append(block(self.inchannel, channels, stride))self.inchannel = channelsreturn nn.Sequential(*layers)def forward(self, x):out = self.conv1(x)out = self.layer1(out)out = self.layer2(out)out = self.layer3(out)out = self.layer4(out)out = F.avg_pool2d(out, 4)out = out.view(out.size(0), -1)out = self.dropout(out)out = self.fc(out)return outclass SparseActivation(nn.Module):act_array=[x.cuda() for x in [nn.ReLU(),nn.ReLU6(),nn.Sigmoid(),nn.Hardsigmoid(),nn.GELU(),nn.SiLU(),nn.Mish(),nn.LeakyReLU(),nn.Hardswish(),nn.PReLU(),nn.SELU(),nn.Softplus(),nn.Softsign()]]def __init__(self,args):super(SparseActivation, self).__init__()self.input_weights = nn.Parameter(torch.randn(1)).cuda()self.act=SparseActivation.act_arrayself.act_weights = nn.Parameter(torch.randn(len(self.act))).cuda()self.args=argsdef forward(self, x):        index=self.args.actif index>=0:index=index-1if index==-1:prob=F.softmax(self.act_weights,dim=0)_, index = torch.topk(prob, 1, dim=0)x=self.act[index](x)if self.args.sparse==0:return xinput=x.flatten(1)input_weights = torch.sigmoid(self.input_weights)        topk = input.size(1)*input_weightstopk=topk.int()topk_vals, topk_indices = torch.topk(input, topk, dim=1)mask = torch.zeros_like(input).scatter(1, topk_indices, topk_vals)return mask.reshape(x.shape)class LitNet(pl.LightningModule):def __init__(self, args):super(LitNet, self).__init__()self.save_hyperparameters()self.args = argsself.resnet18 = ResNet(ResidualBlock)self.criterion = nn.CrossEntropyLoss()self.ws=[]self.replace_activation(self.resnet18,nn.ReLU, SparseActivation,self.ws)    def replace_activation(self,module, old_activation, new_activation,ws):for name, child in module.named_children():if isinstance(child, old_activation):op=new_activation(self.args)ws.append(op.input_weights)setattr(module, name,op)else:self.replace_activation(child, old_activation, new_activation,ws)        def forward(self, x):return self.resnet18(x)def on_train_epoch_start(self):self.train_total_loss=[]self.train_total_acc=[]def on_train_epoch_end(self):self.log('epoch_train_loss', np.mean(self.train_total_loss))self.log('epoch_train_acc', np.mean(self.train_total_acc)) self.log("lr",self.optimizer.state_dict()['param_groups'][0]['lr'])def training_step(self, batch, batch_idx):data, target = batchoutput = self(data)loss = self.criterion(output, target)l2_reg = torch.tensor(0.).cuda()l2_lambda=0.001for param in self.ws:l2_reg += torch.norm(param+4)                    loss += l2_lambda * l2_reg        self.log('iter_train_loss', loss)_, predicted = torch.max(output.data, 1)correct = (predicted == target).sum()acc = 100. * correct / target.size(0)      self.train_total_loss.append(loss.item())self.train_total_acc.append(acc.item())return loss       def on_validation_epoch_start(self):self.val_total_loss=[]self.val_total_acc=[]def on_validation_epoch_end(self):self.log('epoch_val_loss', np.mean(self.val_total_loss))self.log('epoch_val_acc', np.mean(self.val_total_acc))def validation_step(self, batch, batch_idx):data, target = batchoutput = self(data)_, predicted = torch.max(output.data, 1)correct = (predicted == target).sum()acc = 100. * correct / target.size(0)loss = self.criterion(output, target)        self.val_total_loss.append(loss.item())self.val_total_acc.append(acc.item())def test_step(self, batch, batch_idx):data, target = batchoutput = self(data)loss = self.criterion(output, target)self.log('test_loss', loss)return lossdef configure_optimizers(self):self.optimizer = optim.SGD(self.parameters(), lr=self.args.lr, momentum=0.9,weight_decay=5e-4)self.scheduler = optim.lr_scheduler.StepLR(self.optimizer,step_size=10,gamma = 0.8)            return [self.optimizer],[self.scheduler]class CIFAR10DataModule(pl.LightningDataModule):def __init__(self, batch_size):super().__init__()self.batch_size = batch_sizedef setup(self, stage=None):transform_train = transforms.Compose([transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(),transforms.ToTensor(),transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),])transform_test = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),])self.train = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)self.test = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)def train_dataloader(self):return DataLoader(self.train, batch_size=self.batch_size,shuffle=True,num_workers=2,persistent_workers=True)def val_dataloader(self):return DataLoader(self.test, batch_size=self.batch_size,shuffle=False,num_workers=2,persistent_workers=True)def test_dataloader(self):return DataLoader(self.test, batch_size=self.batch_size)def main():parser = argparse.ArgumentParser(description='PyTorch MNIST Example')parser.add_argument('--batch-size', type=int, default=128, metavar='N',help='input batch size for training (default: 64)')parser.add_argument('--epochs', type=int, default=100, metavar='N',help='number of epochs to train (default: 14)')parser.add_argument('--lr', type=float, default=0.01, metavar='LR',help='learning rate (default: 1.0)')parser.add_argument('--act', type=int, default=-1,help='learning rate (default: 1.0)')parser.add_argument('--sparse', type=int, default=0,help='learning rate (default: 1.0)')args = parser.parse_args()cifar10_data = CIFAR10DataModule(batch_size=args.batch_size)log_dir = "lightning_logs"args.sparse=0   #不开启稀疏args.act=0      #自适应激活model = LitNet(args)logger = TensorBoardLogger(save_dir=log_dir, name="SelectiveActive")    trainer = pl.Trainer(logger=logger,devices=1,max_epochs=args.epochs,val_check_interval=1.0,gradient_clip_val=0.9, gradient_clip_algorithm="value")trainer.fit(model, cifar10_data)    args.sparse=0     #不开启稀疏args.act=-1       #不用激活model = LitNet(args)    cifar10_data = CIFAR10DataModule(batch_size=args.batch_size)logger = TensorBoardLogger(save_dir=log_dir, name="NoActive")    trainer = pl.Trainer(logger=logger,devices=1,max_epochs=args.epochs,val_check_interval=1.0,gradient_clip_val=0.9, gradient_clip_algorithm="value")trainer.fit(model, cifar10_data)  args.sparse=1args.act=-1       #不用激活,开启稀疏model = LitNet(args)       logger = TensorBoardLogger(save_dir=log_dir, name="SparseActivation")    trainer = pl.Trainer(logger=logger,devices=1,max_epochs=args.epochs,val_check_interval=1.0,gradient_clip_val=0.9, gradient_clip_algorithm="value")trainer.fit(model, cifar10_data)  for idx,act_name in enumerate(SparseActivation.act_array):name=act_name.__class__.__name__print(name)args.act=idx+1args.sparse=0model = LitNet(args)     logger = TensorBoardLogger(save_dir=log_dir, name=name)    trainer = pl.Trainer(logger=logger,devices=1,max_epochs=args.epochs,val_check_interval=1.0,gradient_clip_val=0.9, gradient_clip_algorithm="value")trainer.fit(model, cifar10_data)if __name__ == '__main__':main()