Alexnet源代码剖析(pytorch版)

@[toc]

环境要求

numpy==1.14.5
Pillow==6.2.0
protobuf==3.6.0
six==1.11.0
tensorboardX==1.2
torch==0.4.0
torchvision==0.2.1

参数配置定义


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# define model parameters
NUM_EPOCHS = 90 # original paper
BATCH_SIZE = 128 # 一次馈入的数据量
MOMENTUM = 0.9 # 优化器的参数
LR_DECAY = 0.0005
LR_INIT = 0.01 #初始学习率
IMAGE_DIM = 227 # pixels;Alexnet model 需要的图片尺寸
NUM_CLASSES = 1000 # 1000 classes for imagenet 2012 dataset
DEVICE_IDS = [0, 1, 2, 3] # GPUs to use;指定多GPU训练的标号
# modify this to point to your data directory
INPUT_ROOT_DIR = 'alexnet_data_in'
TRAIN_IMG_DIR = 'alexnet_data_in/imagenet'
OUTPUT_DIR = 'alexnet_data_out'
LOG_DIR = OUTPUT_DIR + '/tblogs' # tensorboard logs
CHECKPOINT_DIR = OUTPUT_DIR + '/models' # model checkpoints

# make checkpoint path directory
#递归创建文件夹,/home/skzeng/models/model,若目录的models不存在,那会创建models及其下面的model文件夹。对应os.removedirs()
#os.mkdir()则仅能创建一个文件,对应os.rmdir
os.makedirs(CHECKPOINT_DIR, exist_ok=True)

AlexNet网络架构及代码

在这里插入图片描述
Alexnet model class

class AlexNet(nn.Module):
"""
Neural network model consisting of layers propsed by AlexNet paper.
"""
def __init__(self, num_classes=1000):
"""
Define and allocate layers for this neural net.

Args:
num_classes (int): number of classes to predict with this model
"""
#加上该句使得子类能够调用父类的属性
#如果不加,则该子类仅能调用父类的方法
super().__init__()
# input size should be : (b x 3 x 227 x 227)
# The image in the original paper states that width and height are 224 pixels, but
# the dimensions after first convolution layer do not lead to 55 x 55.
#按照一下网络来推理原文中的网络图。通过padding弥补,pooling过的特征图大小,最终维持获得的原文图。
#这里pooling并没有改变特征图的通道数。
#可以将pooling理解为conv layer后面紧贴的一层
# net 定义仅针对数据的操作进行设置,输入输出通道、卷积核,卷积的步数,填充等等;可谓万事具备只差数据,复用就是这样,模子做好,来什么我馈入输出就是了。
self.net = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=96, kernel_size=11, stride=4), # (b x 96 x 55 x 55)
nn.ReLU(),
nn.LocalResponseNorm(size=5, alpha=0.0001, beta=0.75, k=2), # section 3.3
nn.MaxPool2d(kernel_size=3, stride=2), # (b x 96 x 27 x 27)
nn.Conv2d(96, 256, 5, padding=2), # (b x 256 x 27 x 27)
nn.ReLU(),
nn.LocalResponseNorm(size=5, alpha=0.0001, beta=0.75, k=2),
nn.MaxPool2d(kernel_size=3, stride=2), # (b x 256 x 13 x 13)
nn.Conv2d(256, 384, 3, padding=1), # (b x 384 x 13 x 13)
nn.ReLU(),
nn.Conv2d(384, 384, 3, padding=1), # (b x 384 x 13 x 13)
nn.ReLU(),
nn.Conv2d(384, 256, 3, padding=1), # (b x 256 x 13 x 13)
nn.ReLU(),
nn.MaxPool2d(kernel_size=3, stride=2), # (b x 256 x 6 x 6)
)
# classifier is just a name for linear layers
# 全连接层本身就是一个分类器,这里写的好
self.classifier = nn.Sequential(
nn.Dropout(p=0.5, inplace=True),
nn.Linear(in_features=(256 * 6 * 6), out_features=4096),
nn.ReLU(),
nn.Dropout(p=0.5, inplace=True),
nn.Linear(in_features=4096, out_features=4096),
nn.ReLU(),
nn.Linear(in_features=4096, out_features=num_classes),
)
self.init_bias() # initialize bias

#按照原文,初始化权重和偏置
# 遍历nn.Sequential()里面的所有组件
def init_bias(self):
for layer in self.net:
if isinstance(layer, nn.Conv2d):
nn.init.normal_(layer.weight, mean=0, std=0.01)
nn.init.constant_(layer.bias, 0)
# original paper = 1 for Conv2d layers 2nd, 4th, and 5th conv layers
nn.init.constant_(self.net[4].bias, 1)
nn.init.constant_(self.net[10].bias, 1)
nn.init.constant_(self.net[12].bias, 1)

def forward(self, x):
"""
Pass the input through the net.

Args:
x (Tensor): input tensor

Returns:
output (Tensor): output tensor
"""
x = self.net(x)
#类似于reshape()功能
x = x.view(-1, 256 * 6 * 6) # reduce the dimensions for linear layer input
return self.classifier(x)

main.py(data load+train+数据作图记录)

# print the seed value
seed = torch.initial_seed()
print('Used seed : {}'.format(seed))

tbwriter = SummaryWriter(log_dir=LOG_DIR)
print('TensorboardX summary writer created')

# create model,并把模型放入cuda
alexnet = AlexNet(num_classes=NUM_CLASSES).to(device)

# train on multiple GPUs,多GPU计算。
alexnet = torch.nn.parallel.DataParallel(alexnet, device_ids=DEVICE_IDS)
print(alexnet)
print('AlexNet created')

# create dataset and data loader
# transforms.Compose 数据转换组件,内含对数据的一系列操作
# 此处的数据需要具有root/dog/xxx.png的结构,因此需要首先将数据组织成root/类别/图片,的形式
dataset = datasets.ImageFolder(TRAIN_IMG_DIR, transforms.Compose([
# transforms.RandomResizedCrop(IMAGE_DIM, scale=(0.9, 1.0), ratio=(0.9, 1.1)),
transforms.CenterCrop(IMAGE_DIM),
# transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
]))
print('Dataset created')
# 按照参数要求加载数据
dataloader = data.DataLoader(
dataset,
shuffle=True,
pin_memory=True,
num_workers=8,
drop_last=True,
batch_size=BATCH_SIZE)
print('Dataloader created')

# create optimizer
# the one that WORKS
optimizer = optim.Adam(params=alexnet.parameters(), lr=0.0001)
### BELOW is the setting proposed by the original paper - which doesn't train....
# optimizer = optim.SGD(
# params=alexnet.parameters(),
# lr=LR_INIT,
# momentum=MOMENTUM,
# weight_decay=LR_DECAY)
print('Optimizer created')

# multiply LR by 1 / 10 after every 30 epochs,经过一个step_size,按照gamma系数进行缩放学习率
lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)
print('LR Scheduler created')

# start training!!
print('Starting training...')
total_steps = 1
for epoch in range(NUM_EPOCHS):
#每经过一个epoch,下面就自动计数一次,到达阈值则触发衰减学习率
lr_scheduler.step()
for imgs, classes in dataloader:
#将加载到的每一批数据放入cuda
imgs, classes = imgs.to(device), classes.to(device)
output = alexnet(imgs)
loss = F.cross_entropy(output, classes)

# 三部曲,清除冗余梯度;反向传播;优化
optimizer.zero_grad()
loss.backward()
optimizer.step()

# 每10轮计算一次acc
if total_steps % 10 == 0:
with torch.no_grad():
# axis=1则按照横向,返回max的索引和对应的类别
_, pred = torch.max(output, 1)
accuracy = torch.sum(preds == classes)

# log the information and add to tensorboard
# 每10轮做一下可视化图
if total_steps % 10 == 0:
with torch.no_grad():
_, preds = torch.max(output, 1)
accuracy = torch.sum(preds == classes)

print('Epoch: {} \tStep: {} \tLoss: {:.4f} \tAcc: {}'
.format(epoch + 1, total_steps, loss.item(), accuracy.item()))

#记录loss和step之间的关系,一定要用item取tensor里面的具体数值;
#表名:"loss",Y轴:loss.item(),X轴total_steps
tbwriter.add_scalar('loss', loss.item(), total_steps)
tbwriter.add_scalar('accuracy', accuracy.item(), total_steps)

# print out gradient values and parameter average values
if total_steps % 100 == 0:
with torch.no_grad():
# print and save the grad of the parameters
# also print and save parameter values
print('*' * 10)

#每100轮,分别记录网络的 grad和total_steps、weight和total_steps的关系
for name, parameter in alexnet.named_parameters():
if parameter.grad is not None:
avg_grad = torch.mean(parameter.grad)
print('\t{} - grad_avg: {}'.format(name, avg_grad))
tbwriter.add_scalar('grad_avg/{}'.format(name), avg_grad.item(), total_steps)
tbwriter.add_histogram('grad/{}'.format(name),
parameter.grad.cpu().numpy(), total_steps)
if parameter.data is not None:
avg_weight = torch.mean(parameter.data)
print('\t{} - param_avg: {}'.format(name, avg_weight))
tbwriter.add_histogram('weight/{}'.format(name),
parameter.data.cpu().numpy(), total_steps)
tbwriter.add_scalar('weight_avg/{}'.format(name), avg_weight.item(), total_steps)

total_steps += 1

# save checkpoints
checkpoint_path = os.path.join(CHECKPOINT_DIR, 'alexnet_states_e{}.pkl'.format(epoch + 1))
state = {
'epoch': epoch,
'total_steps': total_steps,
'optimizer': optimizer.state_dict(),
'model': alexnet.state_dict(),
'seed': seed,#保留随机种子,以最大可能复现试验
}
torch.save(state, checkpoint_path)

rearrange_voc.py(对原始voc进行修改,以利用)

import os
import os.path as path

VOC_ROOT = path.join('./', 'VOCdevkit', 'VOC2012')
VOC_LABEL_DIR = path.join(VOC_ROOT, 'ImageSets', 'Main')
VOC_IMG_DIR = path.join(VOC_ROOT, 'JPEGImages')
VOC_OUT_REARRANGED = path.join('./', 'voc-data-rearr')

CLASSES = ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']

def move_img(img_name, cls_name):
img_path = path.join(VOC_IMG_DIR, img_name + '.jpg')
out_img_path = path.join(VOC_OUT_REARRANGED, cls_name)
#os的骚操作,cp source-path dist-path
#像这种“{}”.format(),就是把format()内部的内容取代{};简而言之{}是一个数据占位符
os.system('cp {} {}'.format(img_path, out_img_path))

# create output directory
if not os.path.exists(VOC_OUT_REARRANGED):
os.mkdir(VOC_OUT_REARRANGED)

for clsname in CLASSES:
class_folder = path.join(VOC_OUT_REARRANGED, clsname)
if not os.path.exists(class_folder):
os.mkdir(class_folder)

#os.listdir仅获取路径下面一级的文件夹和文件名,返回一个数组
for filename in os.listdir(VOC_LABEL_DIR):
if filename in ['val.txt', 'train.txt', 'trainval.txt']:
continue

# 返回:.之前的内容;.之后的后缀
cls_traintype, _ = path.splitext(filename)
# 根据cls_traintype的内容,以_进行多部分的分割,返回一个数组
clsname, traintype = cls_traintype.split('_')

if traintype != 'train':
continue

print('Processing: {}'.format(filename))

#以对话管理器的形式打开路径文件,执行完毕则自动回收资源

with open(path.join(VOC_LABEL_DIR, filename), 'r') as f:
# 读取文件中的所有行的内容
lines = f.readlines()
for line in lines:
#针对每一行的内容进行strip(),之后进行split()。由于是空,默认为是空格;strip()首先消除首尾的空格,然后使用split()将返回的内容进行分割成多部分

img_name, is_cls = line.strip().split()
#如果是该类别,则将图片进行复制到目标路径
if is_cls == '1':
# 调用移动函数
move_img(img_name, clsname)

 

发表回复

您的电子邮箱地址不会被公开。