1.搭建VGGNet网络
VGG_model.py
# -*-coding:utf-8-*-
import torch.nn as nn
import torch
# official pretrain weights
model_urls = {
'vgg11': 'https://download.pytorch.org/models/vgg11-bbd30ac9.pth',
'vgg13': 'https://download.pytorch.org/models/vgg13-c768596a.pth',
'vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth',
'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth'
}
class VGG(nn.Module):
def __init__(self, features, num_classes=1000, init_weights=False):
super(VGG, self).__init__()
self.features = features
self.classifier = nn.Sequential(
nn.Linear(512*7*7, 4096),
nn.ReLU(True),
nn.Dropout(p=0.5),
nn.Linear(4096, 4096),
nn.ReLU(True),
nn.Dropout(p=0.5),
nn.Linear(4096, num_classes)
)
if init_weights:
self._initialize_weights()
def forward(self, x):
# N x 3 x 224 x 224
x = self.features(x)
# N x 512 x 7 x 7
x = torch.flatten(x, start_dim=1)
# N x 512*7*7
x = self.classifier(x)
return x
def _initialize_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
# nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
nn.init.xavier_uniform_(m.weight)
if m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.Linear):
nn.init.xavier_uniform_(m.weight)
# nn.init.normal_(m.weight, 0, 0.01)
nn.init.constant_(m.bias, 0)
def make_features(cfg: list):
layers = []
in_channels = 3
for v in cfg:
if v == "M":
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
else:
conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
layers += [conv2d, nn.ReLU(True)]
in_channels = v
return nn.Sequential(*layers)
cfgs = {
'vgg11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
'vgg13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
'vgg16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
'vgg19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
}
def vgg(model_name="vgg16", **kwargs):
assert model_name in cfgs, "Warning: model number {} not in cfgs dict!".format(model_name)
cfg = cfgs[model_name]
model = VGG(make_features(cfg), **kwargs)
return model
2.训练网络
VGG_train.py
# -*-coding:utf-8-*-
import os
import json
import torch
import torch.nn as nn
from torchvision import transforms, datasets
import torch.optim as optim
from tqdm import tqdm
from VGG_model import vgg
def main():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("using {} device.".format(device))
data_transform = {
"train": transforms.Compose([transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]),
"val": transforms.Compose([transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])}
data_root = os.path.abspath(os.path.join(os.getcwd(), "../")) # get data root path
image_path = os.path.join(data_root, "data_set") # rock data set path
assert os.path.exists(image_path), "{} path does not exist.".format(image_path)
train_dataset = datasets.ImageFolder(root=os.path.join(image_path, "train"),
transform=data_transform["train"])
train_num = len(train_dataset)
# {'daisy':0, 'dandelion':1, 'roses':2, 'sunflower':3, 'tulips':4}
rock_list = train_dataset.class_to_idx
cla_dict = dict((val, key) for key, val in rock_list.items())
# write dict into json file
json_str = json.dumps(cla_dict, indent=4)
with open('class_indices.json', 'w') as json_file:
json_file.write(json_str)
batch_size = 32
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8]) # number of workers
print('Using {} dataloader workers every process'.format(nw))
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size, shuffle=True,
num_workers=nw)
validate_dataset = datasets.ImageFolder(root=os.path.join(image_path, "val"),
transform=data_transform["val"])
val_num = len(validate_dataset)
validate_loader = torch.utils.data.DataLoader(validate_dataset,
batch_size=batch_size, shuffle=False,
num_workers=nw)
print("using {} images for training, {} images for validation.".format(train_num,
val_num))
# test_data_iter = iter(validate_loader)
# test_image, test_label = test_data_iter.next()
model_name = "vgg16"
net = vgg(model_name=model_name, num_classes=7, init_weights=True)
net.to(device)
loss_function = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.0001)
epochs = 3 #迭代次数
best_acc = 0.0
save_path = './{}Net.pth'.format(model_name)
train_steps = len(train_loader)
for epoch in range(epochs):
# train
net.train()
running_loss = 0.0
train_bar = tqdm(train_loader)
for step, data in enumerate(train_bar):
images, labels = data
optimizer.zero_grad()
outputs = net(images.to(device))
loss = loss_function(outputs, labels.to(device))
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(epoch + 1,
epochs,
loss)
# validate
net.eval()
acc = 0.0 # accumulate accurate number / epoch
with torch.no_grad():
val_bar = tqdm(validate_loader)
for val_data in val_bar:
val_images, val_labels = val_data
outputs = net(val_images.to(device))
predict_y = torch.max(outputs, dim=1)[1]
acc += torch.eq(predict_y, val_labels.to(device)).sum().item()
val_accurate = acc / val_num
print('[epoch %d] train_loss: %.3f val_accuracy: %.3f' %
(epoch + 1, running_loss / train_steps, val_accurate))
if val_accurate > best_acc:
best_acc = val_accurate
torch.save(net.state_dict(), save_path)
print('Finished Training')
if __name__ == '__main__':
main()
3.报错
第一次运行的结果(batch size=32):
内存使用过多导致内存不足,尝试改小batch_size
第二次运行结果(batch size=16):
后期直接调用已经训练好的网络,不需要重新再用数据集训练网络
第三次运行结果(batch_size=5)
/home/xyjin/anaconda3/envs/pytorch/bin/python /home/xyjin/PycharmProjects/data_mining/data_process/rock/classed_data/VGGNet/VGG_train.py
using cpu device.
Using 4 dataloader workers every process
using 12600 images for training, 1400 images for validation.
train epoch[1/3] loss:0.577: 100%|██████████| 2520/2520 [4:22:31<00:00, 6.25s/it]
100%|██████████| 280/280 [09:03<00:00, 1.94s/it]
[epoch 1] train_loss: 1.483 val_accuracy: 0.536
train epoch[2/3] loss:0.895: 52%|█████▏ | 1304/2520 [2:18:56<2:09:24, 6.38s/it]
Process finished with exit code 137 (interrupted by signal 9: SIGKILL)
第四次运行结果(batch size=16)–台式电脑
C:\ProgramData\Anaconda3\envs\pytorch\python.exe C:/Users/pythonProject/data_mining/VGGNet/VGG_train.py
using cpu device.
Using 8 dataloader workers every process
using 12600 images for training, 1400 images for validation.
train epoch[1/3] loss:1.489100%|██████████| 788/788 [3:42:58<00:00, 15.04s/it]
100%|██████████| 88/88 [06:47<00:00, 3.72s/it]
[epoch 1] train_loss: 1.414 val_accuracy: 0.574
train epoch[2/3] loss:0.765100%|██████████| 788/788 [3:44:20<00:00, 15.33s/it]
100%|██████████| 88/88 [06:49<00:00, 3.81s/it]
[epoch 2] train_loss: 1.167 val_accuracy: 0.658
train epoch[3/3] loss:0.837100%|██████████| 788/788 [3:45:21<00:00, 15.00s/it]
100%|██████████| 88/88 [06:44<00:00, 3.76s/it]
[epoch 3] train_loss: 1.039 val_accuracy: 0.691
Finished Training
第五次运行 (迭代10次–台式电脑)
C:\ProgramData\Anaconda3\envs\pytorch\python.exe C:/Users/pythonProject/data_mining/VGGNet/VGG_train.py
using cpu device.
Using 8 dataloader workers every process
using 12600 images for training, 1400 images for validation.
train epoch[1/10] loss:1.395100%|██████████| 788/788 [3:41:22<00:00, 14.86s/it]
100%|██████████| 88/88 [06:46<00:00, 3.77s/it]
[epoch 1] train_loss: 1.445 val_accuracy: 0.515
train epoch[2/10] loss:1.346100%|██████████| 788/788 [3:43:21<00:00, 15.00s/it]
100%|██████████| 88/88 [06:43<00:00, 3.81s/it]
[epoch 2] train_loss: 1.170 val_accuracy: 0.672
train epoch[3/10] loss:1.130100%|██████████| 788/788 [3:44:59<00:00, 14.73s/it]
100%|██████████| 88/88 [06:41<00:00, 3.74s/it]
0%| | 0/788 [00:00<?, ?it/s][epoch 3] train_loss: 1.059 val_accuracy: 0.592
train epoch[4/10] loss:0.653100%|██████████| 788/788 [3:44:20<00:00, 14.96s/it]
100%|██████████| 88/88 [06:42<00:00, 3.73s/it]
0%| | 0/788 [00:00<?, ?it/s][epoch 4] train_loss: 0.967 val_accuracy: 0.660
train epoch[5/10] loss:0.589100%|██████████| 788/788 [3:43:30<00:00, 15.18s/it]
100%|██████████| 88/88 [06:34<00:00, 3.64s/it]
[epoch 5] train_loss: 0.895 val_accuracy: 0.721
train epoch[6/10] loss:1.101100%|██████████| 788/788 [3:44:09<00:00, 14.89s/it]
100%|██████████| 88/88 [06:41<00:00, 3.72s/it]
[epoch 6] train_loss: 0.842 val_accuracy: 0.734
train epoch[7/10] loss:1.600100%|██████████| 788/788 [3:43:06<00:00, 14.81s/it]
100%|██████████| 88/88 [06:39<00:00, 3.72s/it]
[epoch 7] train_loss: 0.785 val_accuracy: 0.760
train epoch[8/10] loss:0.629100%|██████████| 788/788 [3:43:05<00:00, 14.84s/it]
100%|██████████| 88/88 [06:38<00:00, 3.71s/it]
[epoch 8] train_loss: 0.729 val_accuracy: 0.806
train epoch[9/10] loss:0.303100%|██████████| 788/788 [3:42:55<00:00, 14.56s/it]
100%|██████████| 88/88 [06:40<00:00, 3.75s/it]
[epoch 9] train_loss: 0.700 val_accuracy: 0.821
train epoch[10/10] loss:0.748100%|██████████| 788/788 [3:44:18<00:00, 15.19s/it]
100%|██████████| 88/88 [06:46<00:00, 3.75s/it]
[epoch 10] train_loss: 0.655 val_accuracy: 0.829
Finished Training