def train_lenet(split_dir):
"""训练lenet"""
set_seed() # 设置随机种子
rmb_label = {"1": 0, "100": 1}
# 参数设置
MAX_EPOCH = 10
BATCH_SIZE = 16
LR = 0.01
log_interval = 10
val_interval = 1
"""Step 1: 数据读取"""
train_dir = os.path.join(split_dir, "train")
valid_dir = os.path.join(split_dir, "valid")
test_dir = os.path.join(split_dir, "test")
norm_mean = [0.485, 0.456, 0.406]
norm_std = [0.229, 0.224, 0.225]
# 对训练数据进行变换,添加RandomCrop进行数据增强
train_transform = transforms.Compose([
transforms.Resize((32, 32)),
transforms.RandomCrop(32, padding=4),
transforms.ToTensor(),
transforms.Normalize(norm_mean, norm_std),
])
# 对验证数据进行变换
valid_transform = transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize(norm_mean, norm_std),
])
# 构建RMBDataset实例
train_data = RMBDataset(data_dir=train_dir, transform=train_transform)
valid_data = RMBDataset(data_dir=valid_dir, transform=valid_transform)
# 构建DataLoader
train_loader = DataLoader(dataset=train_data,
batch_size=BATCH_SIZE,
shuffle=True)
valid_loader = DataLoader(dataset=valid_data, batch_size=BATCH_SIZE)
"""Step 2: 模型"""
net = LeNet(classes=2)
net.initialize_weights()
"""Step 3: 损失函数"""
criterion = nn.CrossEntropyLoss()
"""Step 4: 优化器"""
optimizer = optim.SGD(net.parameters(), lr=LR, momentum=0.9)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=10,
gamma=0.1)
"""Step 5: 训练"""
train_curve = []
valid_curve = []
figure_dir = os.path.join(os.path.abspath(os.path.dirname(__file__)),
'result')
for epoch in range(MAX_EPOCH):
loss_mean = 0.
correct = 0.
total = 0.
net.train()
for i, data in enumerate(train_loader):
# forward
inputs, labels = data
outputs = net(inputs)
# backward
optimizer.zero_grad()
loss = criterion(outputs, labels)
loss.backward()
# update weights
optimizer.step()
# 统计分类情况
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).squeeze().sum().numpy()
# 打印训练信息
loss_mean += loss.item()
train_curve.append(loss.item())
if (i + 1) % log_interval == 0:
loss_mean = loss_mean / log_interval
print(
"Training:Epoch[{:0>3}/{:0>3}] Iteration[{:0>3}/{:0>3}] Loss: {:.4f} Acc:{:.2%}"
.format(epoch, MAX_EPOCH, i + 1, len(train_loader),
loss_mean, correct / total))
loss_mean = 0.
scheduler.step() # 更新学习率
# validate the model
if (epoch + 1) % val_interval == 0:
correct_val = 0.
total_val = 0.
loss_val = 0.
net.eval()
with torch.no_grad():
for j, data in enumerate(valid_loader):
inputs, labels = data
outputs = net(inputs)
loss = criterion(outputs, labels)
_, predicted = torch.max(outputs.data, 1)
total_val += labels.size(0)
correct_val += (
predicted == labels).squeeze().sum().numpy()
loss_val += loss.item()
valid_curve.append(loss_val / valid_loader.__len__())
print(
"Valid:\t Epoch[{:0>3}/{:0>3}] Iteration[{:0>3}/{:0>3}] Loss: {:.4f} Acc:{:.2%}"
.format(epoch, MAX_EPOCH, j + 1, len(valid_loader),
loss_val, correct_val / total_val))
train_x = range(len(train_curve))
train_y = train_curve
train_iters = len(train_loader)
valid_x = np.arange(
1,
len(valid_curve) + 1
) * train_iters * val_interval # 由于valid中记录的是epochloss，需要对记录点进行转换到iterations
valid_y = valid_curve
plt.plot(train_x, train_y, label='Train')
plt.plot(valid_x, valid_y, label='Valid')
plt.legend(loc='upper right')
plt.ylabel('loss value')
plt.xlabel('Iteration')
figure_path = os.path.join(figure_dir, '0201.png')
plt.savefig(figure_path)
plt.close()
# ============================ step 4/5 优化器 ============================
optimizer = optim.SGD(net.parameters(), lr=LR, momentum=0.9) # 选择优化器
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10,
gamma=0.1) # 设置学习率下降策略
# ============================ step 5/5 训练 ============================
train_curve = list()
valid_curve = list()
for epoch in range(MAX_EPOCH):
loss_mean = 0.
correct = 0.
total = 0.
net.train()
for i, data in enumerate(train_loader):
# forward
inputs, labels = data
outputs = net(inputs)
# backward
optimizer.zero_grad()
loss = criterion(outputs, labels)
loss.backward()
# update weights
optimizer.step()
# 统计分类情况
import os
import random
import cv2
import numpy as np
from lenet import LeNet
def get_training_data(data_dir):
images = []
labels = []
files = os.listdir(data_dir)
random.shuffle(files)
for f in files:
img = cv2.imread(os.path.join(data_dir, f), cv2.IMREAD_GRAYSCALE)
img = cv2.resize(img, (32, 32))
img = img.astype(np.float32).reshape(32, 32, 1) / 255.0
images.append(img)
num = int(f[0])
label = np.zeros(10, dtype=np.float32)
label[num] = 1
labels.append(label)
return (np.array(images), np.array(labels))
if __name__ == '__main__':
x, y = get_training_data("mnist/train")
lenet = LeNet()
lenet.train(x, y)
lenet.save("lenet.npy")
batch_size=args.batch_size,
shuffle=True,
pin_memory=cuda)
testloader = torch.utils.data.DataLoader(testset,
batch_size=args.batch_size,
shuffle=False,
pin_memory=cuda)
model = LeNet(k=len(classes), args=args)
if cuda:
model.cuda()
for epoch in range(1, args.epochs + 1):
model.train()
for i, (x_batch, y_batch) in enumerate(trainloader):
if cuda:
x_batch, y_batch = x_batch.cuda(), y_batch.cuda()
model.optimizer.zero_grad()
# forward + backward + optimize
output = model(x_batch)
loss = model.criterion(output, y_batch)
loss.backward()
model.optimizer.step()
if i % args.log_interval == 0:
print('Epoch: {:3d}, Batch {:3d}/{}, Loss: {:.5f}'.format(
class Ui(QtWidgets.QMainWindow):
buttons = [
"load_image", "color_conversion", "image_flipping",
"blending", "global_threshold", "local_threshold",
"gaussian", "sobel_x", "sobel_y", "magnitude", "rst",
"show_train_image", "show_hyper", "train_1", "pt",
"inference", "ok", "show_train_result", "cancel"]
inputs = ["angle", "scale", "tx", "ty", "test_index"]
def __init__(self):
super(Ui, self).__init__()
uic.loadUi('main_window.ui', self)
self.get_widgets()
self.get_input()
self.bind_event()
self.param_setup()
self.torch_setup()
self.show()
def get_widgets(self):
for btn in self.buttons:
setattr(self, btn, self.findChild(QtWidgets.QPushButton, btn))
def get_input(self):
for inp in self.inputs:
setattr(self, inp, self.findChild(QtWidgets.QLineEdit, inp))
def bind_event(self):
for btn in self.buttons:
getattr(self, btn).clicked.connect(partial(
getattr(events, btn),
self))
def param_setup(self):
self.batch_size = 32
self.learning_rate = 0.001
self.opt = "SGD"
self.loss_list = []
self.loss_epoch = []
self.acc_train_epoch = []
self.acc_test_epoch = []
self.compose = transforms.Compose([
transforms.Resize((32,32)),
transforms.ToTensor()
])
def torch_setup(self):
self.data_train = MNIST('./data/mnist',
train=True,
download=True,
transform=self.compose)
self.data_test = MNIST('./data/mnist',
train=False,
download=True,
transform=self.compose)
self.data_train_loader = DataLoader(self.data_train, batch_size=self.batch_size, shuffle=True, num_workers=4)
self.data_test_loader = DataLoader(self.data_test, batch_size=self.batch_size, num_workers=4)
self.criterion = nn.CrossEntropyLoss()
self.net = LeNet()
self.optimizer = getattr(optim, self.opt)(self.net.parameters(), lr=self.learning_rate)
try:
self.net.load_state_dict(load('model_params.pkl'))
self.loaded = True
print("Loaded")
except Exception as e:
print(e)
self.loaded = False
print("Not loaded")
def train(self, epoch):
self.net.train()
self.loss_list = []
correct, total = 0, 0
for i, (images, labels) in enumerate(self.data_train_loader):
self.optimizer.zero_grad()
output = self.net(images)
loss = self.criterion(output, labels)
pred = output.data.max(1, keepdim=True)[1]
correct += np.sum(np.squeeze(pred.eq(labels.data.view_as(pred))).cpu().numpy())
total += images.size(0)
self.loss_list.append(loss.detach().cpu().item())
if i % 100 == 0:
print(f'Train - Epoch {epoch}, Batch: {i}, Loss: {loss.detach().cpu().item()}')
loss.backward()
self.optimizer.step()
self.acc_train_epoch.append(correct/total)
self.loss_epoch.append(sum(self.loss_list)/len(self.loss_list))
def test(self):
self.net.eval()
total_correct, avg_loss = 0, 0.0
for i, (images, labels) in enumerate(self.data_test_loader):
output = self.net(images)
avg_loss += self.criterion(output, labels).sum()
pred = output.detach().max(1)[1]
total_correct += pred.eq(labels.view_as(pred)).sum()
avg_loss /= len(self.data_test)
acc = float(total_correct)/len(self.data_test)
self.acc_test_epoch.append(acc)
def test_and_train(self, epoch):
self.train(epoch)
self.test()
def training(model_name, trainloader, validloader, input_channel=3, epochs=1, resume=True, self_define=True, only_print=False):
# load self defined or official net
assert model_name in ["LeNet", "VGG16", "ResNet", "DenseNet"]
if self_define:
if model_name == "LeNet":
net = LeNet(input_channel)
elif model_name == "VGG16":
net = VGG16(input_channel)
elif model_name == "ResNet":
net = ResNet(input_channel)
elif model_name == "DenseNet":
net = DenseNet(input_channel)
else:
if model_name == "LeNet":
net = LeNet(input_channel) # on official LeNet
elif model_name == "VGG16":
net = models.vgg16_bn(pretrained=False, num_classes=10)
elif model_name == "ResNet":
net = models.resnet50(pretrained=False, num_classes=10)
elif model_name == "DenseNet":
net = models.DenseNet(num_classes=10)
# sum of net parameters number
print("Number of trainable parameters in %s : %f" % (model_name, sum(p.numel() for p in net.parameters() if p.requires_grad)))
# print model structure
if only_print:
print(net)
return
# resume training
param_path = "./model/%s_%s_parameter.pt" % (model_name, "define" if self_define else "official")
if resume:
if os.path.exists(param_path):
net.load_state_dict(torch.load(param_path))
net.train()
print("Resume training " + model_name)
else:
print("Train %s from scratch" % model_name)
else:
print("Train %s from scratch" % model_name)
# define loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
# train on GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print('train on %s' % device)
net.to(device)
running_loss = 0.0
train_losses = []
valid_losses = []
mini_batches = 125 * 5
for epoch in range(epochs):
for i, data in enumerate(trainloader, 0):
# get one batch
# inputs, labels = data
inputs, labels = data[0].to(device), data[1].to(device)
# switch model to training mode, clear gradient accumulators
net.train()
optimizer.zero_grad()
# forward + backward + optimize
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
if i % mini_batches == mini_batches - 1: # print and valid every <mini_batches> mini-batches
# validate model in validation dataset
valid_loss = valid(net, validloader, criterion, device)
print('[%d, %5d] train loss: %.3f, validset loss: %.3f' % (
epoch + 1, i + 1, running_loss / mini_batches, valid_loss))
train_losses.append(running_loss / mini_batches)
valid_losses.append(valid_loss)
running_loss = 0.0
# save parameters
torch.save(net.state_dict(), param_path)
# # save checkpoint
# torch.save({
# 'epoch': epoch,
# 'model_state_dict': net.state_dict(),
# 'optimizer_state_dict': optimizer.state_dict(),
# 'loss': loss
# }, "./checkpoints/epoch_" + str(epoch) + ".tar")
print('Finished Training, %d images in all' % (len(train_losses) * batch_size * mini_batches / epochs))
# draw loss curve
assert len(train_losses) == len(valid_losses)
loss_x = range(0, len(train_losses))
plt.plot(loss_x, train_losses, label="train loss")
plt.plot(loss_x, valid_losses, label="valid loss")
plt.title("Loss for every %d mini-batch" % mini_batches)
plt.xlabel("%d mini-batches" % mini_batches)
plt.ylabel("Loss")
plt.legend()
plt.savefig(model_name + "_loss.png")
plt.show()
train_generator = Datagenerator(file, 'train', num, batch_size)
test_generator = Datagenerator(file, 'test', num, 200)
#define and initialize LeNet-5
net = LeNet(rbf_w)
net.init_weights()
step = 0
for eps in range(epochs):
train_generator.get_data('batch', True)
while True:
try:
train_images, train_labels = train_generator.gen_batch()
train_loss, train_acc = net.train(train_images, train_labels)
if step % 10 == 0:
record_loss.append(train_loss)
record_acc.append(train_acc)
print('step --> %d, loss : %.4f, acc : %.2f' %
(step, train_loss, train_acc))
step += 1
except OutOfRange:
# finish a epoch, validate test set
test_generator.get_data('batch', False)
tacc = []
while True:
try:
test_images, test_labels = test_generator.gen_batch()
