rock数据集VGGNet实现

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