#!/usr/bin/env python
# -*- coding:utf-8-*-
# file: test_mnist_centerloss.py
# @author: jory.d
# @contact: dangxusheng163@163.com
# @time: 2019/11/29 11:11
# @desc:
"""
conda install pytorch=1.11 torchvision=0.12 -c conda-forge
python 3.8
pytorch=1.11.0=cuda112py38habe9d5a_202
torchvision=0.12.0=cuda112py38h46b2766_1
cudatoolkit=11.8
cudnn=8.9.7
"""
import os
import os.path as osp
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = str(0)
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as TF
import torchvision.datasets as datasets
import torch.optim.lr_scheduler as lr_scheduler
from torch.utils.data import DataLoader
import cv2
batch_size = 128
base_lr = 0.005
test_batch = 256
weight_decay = 0.0001
epochs = 10
gradient_clip = 5
num_classes = 10
input_size = 56
input_ch = 1
last_epoch = -1
eval_freq = 2
dataset_root = '/home/dangxs/projects/ai_dataset'
save_model_path = './mnist_test/mnist_classifer.pth'
onnx_model_path = './mnist_test/mnist_classifer.onnx'
os.makedirs(osp.dirname(save_model_path), exist_ok=True)
def get_mnist_loader(transform, is_train=True):
"""
mnist 数据集加载
目录结构如下：
-- a
--MNIST_Data
--MNIST
--processed
--training.pt
--test.pt
--raw
--t10k-images-idx3-ubyte
--t10k-labels-idx1-ubyte
--train-images-idx3-ubyte
--train-labels-idx1-ubyte
"""
train_loader = DataLoader(
datasets.MNIST(
# ../MNIST_Data/mnist/MNIST/processed/training.pt
# ../MNIST_Data/mnist/MNIST/processed/raw/t10k-images-idx3-ubyte
dataset_root + "/mnist/MNIST_Data",
# "../MNIST_Data/mnist",
train=is_train,
download=False,
transform=transform,
),
batch_size=batch_size if is_train else test_batch,
shuffle=True,
num_workers=4, pin_memory=True
)
return train_loader
def get_cifar10_loader(transform, is_train=True):
train_loader = DataLoader(
datasets.CIFAR10(
# ../CIFAR10_Data/cifar10/cifar-10-batches-py/data_batch_1
# ../CIFAR10_Data/cifar10/cifar-10-batches-py/data_batch_2
dataset_root + "/CIFAR10_Data/cifar10",
train=is_train,
download=False,
transform=transform,
),
batch_size=batch_size if is_train else test_batch,
shuffle=True,
num_workers=4, pin_memory=True
)
return train_loader
class Classifier(nn.Module):
def __init__(self, in_c, in_size, n_classes):
super(Classifier, self).__init__()
self.layers = nn.Sequential(
nn.Conv2d(in_c, 16, 3, 1, 1),
nn.ReLU(True),
nn.MaxPool2d(2, 2), # 1/2
nn.Conv2d(16, 32, 3, 1, 1),
nn.ReLU(True),
nn.MaxPool2d(2, 2), # 1/2
nn.Conv2d(32, 64, 3, 1, 1),
nn.ReLU(True),
nn.MaxPool2d(2, 2), # 1/2
)
scale = 2 ** 3
out_conv_size = in_size // scale
self.fc = nn.Linear(64 * out_conv_size * out_conv_size, n_classes)
def forward(self, x):
x = self.layers(x)
# print("x.shape: ", x.shape)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
def get_model():
# net = minst_model.Classifier(num_classes)
# net = minst_model_bak.Classifier(num_classes)
# net = reneXt_Net2_mnist.Classifier(input_ch, input_size, num_classes)
# net = resneXt_Net2.Classifier(input_ch, input_size, num_classes)
# net = reneXt_Net2.Classifier(input_ch, input_size, num_classes)
net = Classifier(1, input_size, num_classes)
return net
def get_current_lr(optimizer):
for param_group in optimizer.param_groups:
return param_group['lr']
def get_model_parameters(model):
"""
获取模型的总参数, 可训练参数, 模型大小MB
:param model:
:return:
"""
total_params = sum(param.numel() for param in model.parameters())
train_params = sum(param.numel() for param in model.parameters() if param.requires_grad)
model_size = total_params * 4 / 1024 / 1024
return total_params, train_params, model_size
def get_scheduler(optimizer):
lr_policy = 'linear'
if lr_policy == 'linear':
start_epoch = 5 # 从 5 epoch 开始衰减
niter_decay = 2 # 每隔 2 epoch 衰减一次
def lambda_rule(epoch):
lr_l = 1.0 - max(0, epoch - start_epoch) / float(niter_decay + 1)
lr_l = max(0.000001, lr_l)
return lr_l
scheduler = lr_scheduler.LambdaLR(optimizer, lambda_rule)
elif lr_policy == 'step':
decay_step = 30 # 每过 30 epoch 衰减一次
decay_factor = 0.1 # 衰减因子，
# 当epoch每过stop_size时, 学习率都变为初始学习率的gamma倍
# Assuming optimizer uses lr = 0.05 for all groups
# lr = 0.05 if epoch < 30
# lr = 0.005 if 30 <= epoch < 60
# lr = 0.0005 if 60 <= epoch < 90
scheduler = lr_scheduler.StepLR(optimizer, step_size=decay_step, gamma=decay_factor)
elif lr_policy == 'exp':
decay_factor = 0.1 # 衰减因子，
scheduler = lr_scheduler.ExponentialLR(optimizer, decay_factor, last_epoch=-1)
elif lr_policy == 'plateau':
# 当网络的评价指标不在提升的时候，可以通过降低网络的学习率来提高网络性能
# patience : 容忍网路的性能不提升的次数，高于这个次数就降低学习率
# factor : 学习率每次降低多少，new_lr = old_lr * factor
# threshold : 测量新最佳值的阈值，仅关注重大变化。
# min_lr,学习率的下限
scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, min_lr=0.000001, mode='min',
factor=0.2, threshold=0.01, patience=5)
elif lr_policy == 'cosine':
# 余弦退火学习率策略
"""
以余弦函数为周期，并在每个周期最大值时重新设置学习率。以初始学习率为最大学习率，
以 2∗Tmax2∗Tmax 为周期，在一个周期内先下降，后上升。
T_max(int) : 一次学习率周期的迭代次数，即 T_max 个 epoch 之后重新设置学习率。
eta_min : 最小学习率，即在一个周期中，学习率最小会下降到 eta_min，默认值为 0。
"""
scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=10, eta_min=0.000001)
else:
return NotImplementedError('learning rate policy [%s] is not implemented', lr_policy)
return scheduler
def main():
global last_epoch
# define netowrk
net = get_model().cuda()
total_params, train_params, model_size = get_model_parameters(net)
print(" == total parameters: " + str(total_params))
print(" == train parameters: " + str(train_params))
print(" == total model size: %s MB" % model_size)
criterion = nn.CrossEntropyLoss().cuda() # softmax loss
optimizer = torch.optim.SGD(net.parameters(),
lr=base_lr,
momentum=0.9,
weight_decay=weight_decay,
nesterov=True)
# define lr policy
lr_adjust_scheduler = get_scheduler(optimizer)
# load training data, do data augmentation and get data loader
transform_train = TF.Compose([
TF.Resize((input_size,input_size)),
TF.Grayscale(1),
# # ================================================================================================
# TF.TenCrop(28), # this is a list of PIL Images
# TF.Lambda(lambda crops: torch.stack([TF.ToTensor()(crop) for crop in crops])), # returns a 4D tensor
# TF.Lambda(lambda crops: torch.stack([TF.Normalize([0.5], [0.5])(crop) for crop in crops]))
# # ================================================================================================
TF.ToTensor(),
TF.Normalize([0.5], [0.5])
])
transform_test = TF.Compose([
TF.Resize((input_size, input_size)),
TF.Grayscale(1),
TF.ToTensor(),
TF.Normalize([0.5], [0.5])
])
train_loader = get_mnist_loader(transform_train, True)
test_loader = get_mnist_loader(transform_test, False)
# train_loader = get_cifar10_loader(transform_train, True)
# test_loader = get_cifar10_loader(transform_test, False)
print(" ======= Training =======\n")
for epoch in range(last_epoch + 1, epochs):
train(train_loader, net, criterion, optimizer, epoch)
if epoch == 0 or (epoch + 1) % eval_freq == 0 or epoch == epochs - 1:
validate(test_loader, net, criterion, optimizer, epoch)
lr_adjust_scheduler.step()
torch.cuda.empty_cache()
def train(train_loader, net, criterion, optimizer, epoch):
start = time.time()
net.train()
train_loss, correct, total = 0, 0, 0
print(" === Epoch: [{}/{}] === ".format(epoch + 1, epochs))
for batch_index, (inputs, targets) in enumerate(train_loader):
# move tensor to GPU
inputs, targets = inputs.cuda(), targets.cuda()
# # ==============================================================
# # batchs_size 变成原来的10倍。label也要做处理
# this_batch_size = targets.size(0)
# bs, ncrops, c, h, w = inputs.size() # input is a 5d tensor, target is 2d
# inputs = inputs.view(-1, c, h, w)
# targets1 = torch.zeros([inputs.size(0)]).long().cuda()
# for i in range(this_batch_size):
# targets1[i * 10:(i + 1) * 10] = targets[i]
#
# # shuffle x & y
# shuffle_idx = torch.randperm(targets1.size(0))
# inputs = inputs[shuffle_idx]
# targets = targets1[shuffle_idx]
# # ==============================================================
outputs = net(inputs)
loss = criterion(outputs, targets)
# zero the gradient buffers
optimizer.zero_grad()
# backward
loss.backward()
# clip gradient
nn.utils.clip_grad_norm_(net.parameters(), gradient_clip)
# update weight
optimizer.step()
# count the loss and acc
train_loss += loss.item()
predicted = torch.argmax(outputs, dim=-1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
if (batch_index + 1) % 50 == 0:
print(" == step: [{:3}/{}], train loss: {:.3f} | train acc: {:6.3f}% | lr: {:.6f}".format(
batch_index + 1, len(train_loader),
train_loss / (batch_index + 1), 100.0 * correct / total, get_current_lr(optimizer)))
end = time.time()
print(" == cost time: {:.4f}s".format(end - start))
train_loss = train_loss / (batch_index + 1)
train_acc = correct / total
return train_loss, train_acc
def validate(test_loader, net, criterion, optimizer, epoch):
global best_prec
net.eval()
test_loss, correct, total = 0, 0, 0
print(" === Validate ===".format(epoch + 1, epochs))
with torch.no_grad():
for batch_index, (inputs, targets) in enumerate(test_loader):
inputs, targets = inputs.cuda(), targets.cuda()
outputs = net(inputs)
loss = criterion(outputs, targets)
probs = F.softmax(outputs, dim=-1)
test_loss += loss.item()
predicted = torch.argmax(probs, dim=-1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
# save checkpoint
state = {
'state_dict': net.state_dict(),
'last_epoch': epoch,
'optimizer': optimizer.state_dict(),
}
torch.save(state, save_model_path)
print(" == test loss: {:.3f} | test acc: {:6.3f}%".format(
test_loss / (batch_index + 1), 100.0 * correct / total))
test_loss = test_loss / (batch_index + 1)
test_acc = correct / total
return test_loss, test_acc
def export_onnx():
model = get_model().cuda()
assert osp.exists(save_model_path)
checkpoint = torch.load(save_model_path)
model.load_state_dict(checkpoint['state_dict'], strict=True)
model.eval()
dummy_input = torch.randn(1, input_ch, input_size, input_size).cuda()
torch.onnx.export(model, dummy_input, onnx_model_path)
from PIL import Image
def print_mid_layers():
model = get_model().cuda()
assert osp.exists(save_model_path)
checkpoint = torch.load(save_model_path)
model.load_state_dict(checkpoint['state_dict'], strict=True)
model.eval()
print('loaded is done.')
root = '/home/lpadas1/share/HDD/jory.d/dataset'
image_path = root + '/face_expression_dataset/others/download_faces_nolabels/17-0_1.bmp'
assert osp.exists(image_path)
transform_test = TF.Compose([
TF.Resize((input_size, input_size)),
TF.Grayscale(1),
TF.ToTensor(),
TF.Normalize([0.5], [0.5]),
])
image_tensor = transform_test(Image.open(image_path).convert('RGB'))
# # print(image_tensor.size())
image_tensor = torch.unsqueeze(image_tensor, 0).cuda().float()
def hook(module, input, output):
md_output = output.clone()
md_output = torch.flatten(md_output)[:10]
print(f'{module}: {md_output}')
# 注册hook function
for name, md in model.named_modules():
if name == '': continue
if len(md._modules) > 0:
continue
else:
md.register_forward_hook(hook)
with torch.no_grad():
outputs = model(image_tensor)
probs = torch.nn.functional.softmax(outputs, dim=-1)
predicted = torch.argmax(probs, dim=-1)
print('fc output: ', outputs)
print('softmax output: ', probs)
print('pred: ', predicted)
if __name__ == '__main__':
# print_mid_layers()
# export_onnx()
main()
# net = Classifier(1,28,10)
# input_shape = (1, 28, 28)
# input = torch.randn(1, *input_shape)
# out = net(input)
# print(out)
浙公网安备 33010602011771号