diff --git a/data_module.py b/data_module.py index 831a695..2069df2 100644 --- a/data_module.py +++ b/data_module.py @@ -1,9 +1,11 @@ +import glob import os -import numpy +import random import torch -from torch import Tensor from torch.utils.data import Dataset, DataLoader import pytorch_lightning as pl +from PIL import Image +from torchvision import transforms class DataModule(pl.LightningDataModule): @@ -17,47 +19,45 @@ class DataModule(pl.LightningDataModule): self.dataset_path = dataset_path def setup(self, stage=None) -> None: - # 得到全部数据的list - dataset_list = self.get_dataset_list() + k_fold_dataset_list = self.get_k_fold_dataset_list() if stage == 'fit' or stage is None: - dataset_train, dataset_val = self.get_dataset_lists(dataset_list) - self.train_dataset = CustomDataset(dataset_train, self.config) - self.val_dataset = CustomDataset(dataset_val, self.config) + dataset_train, dataset_val = self.get_fit_dataset_lists(k_fold_dataset_list) + self.train_dataset = CustomDataset(self.dataset_path, dataset_train, self.config, 'train') + self.val_dataset = CustomDataset(self.dataset_path, dataset_val, self.config, 'train') if stage == 'test' or stage is None: - self.test_dataset = CustomDataset(dataset_list, self.config) + dataset_test = self.get_test_dataset_lists(k_fold_dataset_list) + self.test_dataset = CustomDataset(self.dataset_path, dataset_test, self.config, 'test') - def get_dataset_list(self): - if not os.path.exists(self.dataset_path + '/dataset_list.txt'): - # 针对数据拟合获得dataset - dataset = torch.randn(self.config['dataset_len'], self.config['dim_in'] + 1) - noise = torch.randn(self.config['dataset_len']) - dataset[:, self.config['dim_in']] = torch.cos(1.5 * dataset[:, 0]) * (dataset[:, 1] ** 2.0) + torch.cos( - torch.sin(dataset[:, 2] ** 3)) + torch.arctan(dataset[:, 4]) + noise - assert (dataset[torch.isnan(dataset)].shape[0] == 0) - written = [' '.join([str(temp) for temp in dataset[cou, :].tolist()]) for cou in range(dataset.shape[0])] + def get_k_fold_dataset_list(self): + # 得到用于K折分割的数据的list, 并生成文件夹进行保存 + if not os.path.exists(self.dataset_path + '/k_fold_dataset_list.txt'): + # 获得用于k折分割的数据的list + dataset = glob.glob(self.dataset_path + '/train/image/*.png') + random.shuffle(dataset) + written = dataset - with open(self.dataset_path + '/dataset_list.txt', 'w', encoding='utf-8') as f: + with open(self.dataset_path + '/k_fold_dataset_list.txt', 'w', encoding='utf-8') as f: for line in written: - f.write(line + '\n') - print('已生成新的数据list') + f.write(line.replace('\\', '/') + '\n') + print('已生成新的k折数据list') else: - dataset_list = open(self.dataset_path + '/dataset_list.txt').readlines() - # 针对数据拟合获得dataset - dataset_list = [[float(temp) for temp in item.strip('\n').split(' ')] for item in dataset_list] - dataset = torch.Tensor(dataset_list).float() - + dataset = open(self.dataset_path + '/k_fold_dataset_list.txt').readlines() + dataset = [item.strip('\n') for item in dataset] return dataset - def get_dataset_lists(self, dataset_list: Tensor): + def get_fit_dataset_lists(self, dataset_list: list): # 得到一个fold的数据量和不够组成一个fold的剩余数据的数据量 - num_1fold, remainder = divmod(self.config['dataset_len'], self.k_fold) + num_1fold, remainder = divmod(len(dataset_list), self.k_fold) # 分割全部数据, 得到训练集, 验证集, 测试集 - dataset_val = dataset_list[num_1fold * self.kth_fold:(num_1fold * (self.kth_fold + 1) + remainder), :] - temp = torch.ones(dataset_list.shape[0]) - temp[num_1fold * self.kth_fold:(num_1fold * (self.kth_fold + 1) + remainder)] = 0 - dataset_train = dataset_list[temp == 1] + dataset_val = dataset_list[num_1fold * self.kth_fold:(num_1fold * (self.kth_fold + 1) + remainder)] + del (dataset_list[num_1fold * self.kth_fold:(num_1fold * (self.kth_fold + 1) + remainder)]) + dataset_train = dataset_list return dataset_train, dataset_val + def get_test_dataset_lists(self, dataset_list): + dataset = glob.glob(self.dataset_path + '/test/image/*.png') + return dataset + def train_dataloader(self): return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=self.num_workers, pin_memory=True) @@ -72,13 +72,19 @@ class DataModule(pl.LightningDataModule): class CustomDataset(Dataset): - def __init__(self, dataset, config): + def __init__(self, dataset_path, dataset, config, type): super().__init__() - self.x = dataset[:, 0:config['dim_in']] - self.y = dataset[:, config['dim_in']] + self.dataset = dataset + self.trans = transforms.ToTensor() + self.labels = open(dataset_path + '/' + type + '/label.txt').readlines() def __getitem__(self, idx): - return self.x[idx, :], self.y[idx] + image_path = self.dataset[idx] + image_name = os.path.basename(image_path) + image = Image.open(image_path) + image = self.trans(image) + label = torch.Tensor([int(self.labels[int(image_name.strip('.png'))].strip('\n'))]) + return image_name, image, label.long() def __len__(self): - return self.x.shape[0] + return len(self.dataset) diff --git a/main.py b/main.py index 7a24c35..e6c8ce8 100644 --- a/main.py +++ b/main.py @@ -32,15 +32,15 @@ def main(stage, :param precision: 训练精度, 正常精度为32, 半精度为16, 也可以是64. 精度代表每个参数的类型所占的位数 :param seed: - :param dataset_path: 数据集地址, 其目录下包含数据集, 标签, 全部数据的命名list + :param dataset_path: 数据集地址, 其目录下包含数据集文件夹, 标签文件夹, 全部数据的命名list :param gpus: :param tpu_cores: - :param version_nth: 该folds的第一个版本的版本号 + :param version_nth: 不论是重载训练还是测试, 固定为该folds的第一个版本的版本号 :param path_final_save: :param every_n_epochs: 每n个epoch设置一个检查点 :param save_top_k: :param kth_fold_start: 从第几个fold开始, 若使用重载训练, 则kth_fold_start为重载第几个fold, 第一个值为0. - 非重载训练的情况下, 可以通过调整该值控制训练的次数 + 非重载训练的情况下, 可以通过调整该值控制训练的次数; :param k_fold: """ # 经常改动的 参数 作为main的输入参数 @@ -70,8 +70,8 @@ def main(stage, save_checkpoint = SaveCheckpoint(seed=seed, max_epochs=max_epochs, path_final_save=path_final_save, every_n_epochs=every_n_epochs, verbose=True, - monitor='Validation loss', save_top_k=save_top_k, - mode='min') + monitor='Validation acc', save_top_k=save_top_k, + mode='max') training_module = TrainModule(config=config) if kth_fold != kth_fold_start or load_checkpoint_path is None: print('进行初始训练') @@ -101,6 +101,6 @@ def main(stage, if __name__ == "__main__": main('fit', max_epochs=2, batch_size=32, precision=16, seed=1234, dataset_path='./dataset', k_fold=5, # gpus=1, - # version_nth=8, # 该folds的第一个版本的版本号 + # version_nth=8, kth_fold_start=4, ) diff --git a/network/MLP_JDLU.py b/network/MLP_JDLU.py deleted file mode 100644 index fe4baa7..0000000 --- a/network/MLP_JDLU.py +++ /dev/null @@ -1,51 +0,0 @@ -import math - -import torch.nn as nn -from network_module.activation import jdlu, JDLU - - -class MLPLayer(nn.Module): - def __init__(self, dim_in, dim_out, res_coef=0.0, dropout_p=0.1): - super().__init__() - self.linear = nn.Linear(dim_in, dim_out) - self.res_coef = res_coef - self.activation = JDLU(dim_out) - self.dropout = nn.Dropout(dropout_p) - self.ln = nn.LayerNorm(dim_out) - - def forward(self, x): - y = self.linear(x) - y = self.activation(y) - y = self.dropout(y) - if self.res_coef == 0: - return y - else: - return self.res_coef * x + y - - -class MLP_JDLU(nn.Module): - def __init__(self, dim_in, dim, res_coef=0.5, dropout_p=0.1, n_layers=10): - super().__init__() - self.mlp = nn.ModuleList() - self.first_linear = MLPLayer(dim_in, dim) - self.n_layers = n_layers - for i in range(n_layers): - self.mlp.append(MLPLayer(dim, dim, res_coef, dropout_p)) - self.final = nn.Linear(dim, 1) - self.apply(self.weight_init) - - def forward(self, x): - x = self.first_linear(x) - for layer in self.mlp: - x = layer(x) - x = self.final(x) - return x.squeeze() - - @staticmethod - def weight_init(m): - if isinstance(m, nn.Linear): - nn.init.xavier_normal_(m.weight) - - fan_in, _ = nn.init._calculate_fan_in_and_fan_out(m.weight) - bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 - nn.init.uniform_(m.bias, -bound, bound) diff --git a/network/MLP_ReLU.py b/network/MLP_ReLU.py deleted file mode 100644 index e9d67fe..0000000 --- a/network/MLP_ReLU.py +++ /dev/null @@ -1,41 +0,0 @@ -import torch.nn as nn -from network_module.activation import jdlu, JDLU - - -class MLPLayer(nn.Module): - def __init__(self, dim_in, dim_out, res_coef=0.0, dropout_p=0.1): - super().__init__() - self.linear = nn.Linear(dim_in, dim_out) - self.res_coef = res_coef - self.activation = nn.ReLU() - self.activation1 = JDLU(dim_out) - self.dropout = nn.Dropout(dropout_p) - self.ln = nn.LayerNorm(dim_out) - - def forward(self, x): - y = self.linear(x) - y = self.activation1(y) - # y = jdlu(y) - y = self.dropout(y) - if self.res_coef == 0: - return self.ln(y) - else: - return self.ln(self.res_coef * x + y) - - -class MLP(nn.Module): - def __init__(self, dim_in, dim, res_coef=0.5, dropout_p=0.1, n_layers=10): - super().__init__() - self.mlp = nn.ModuleList() - self.first_linear = MLPLayer(dim_in, dim) - self.n_layers = n_layers - for i in range(n_layers): - self.mlp.append(MLPLayer(dim, dim, res_coef, dropout_p)) - self.final = nn.Linear(dim, 1) - - def forward(self, x): - x = self.first_linear(x) - for layer in self.mlp: - x = layer(x) - x = self.final(x) - return x.squeeze() diff --git a/network/res_net.py b/network/res_net.py new file mode 100644 index 0000000..1f62620 --- /dev/null +++ b/network/res_net.py @@ -0,0 +1,185 @@ +''' +Properly implemented ResNet-s for CIFAR10 as described in paper [1]. + +The implementation and structure of this file is hugely influenced by [2] +which is implemented for ImageNet and doesn't have option A for identity. +Moreover, most of the implementations on the web is copy-paste from +torchvision's resnet and has wrong number of params. + +Proper ResNet-s for CIFAR10 (for fair comparision and etc.) has following +number of layers and parameters: + +name | layers | params +ResNet20 | 20 | 0.27M +ResNet32 | 32 | 0.46M +ResNet44 | 44 | 0.66M +ResNet56 | 56 | 0.85M +ResNet110 | 110 | 1.7M +ResNet1202| 1202 | 19.4m + +which this implementation indeed has. + +Reference: +[1] Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun + Deep Residual Learning for Image Recognition. arXiv:1512.03385 +[2] https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py + +If you use this implementation in you work, please don't forget to mention the +author, Yerlan Idelbayev. +''' +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.nn.init as init + +from torch.autograd import Variable + +__all__ = ['ResNet', 'resnet20', 'resnet32', 'resnet44', 'resnet56', 'resnet110', 'resnet1202'] + + +def _weights_init(m): + if isinstance(m, nn.Linear): + nn.init.kaiming_normal_(m.weight) + + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(m.weight) + bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 + nn.init.uniform_(m.bias, -bound, bound) + + elif isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight) + + +class LambdaLayer(nn.Module): + def __init__(self, lambd): + super(LambdaLayer, self).__init__() + self.lambd = lambd + + def forward(self, x): + return self.lambd(x) + + +class BasicBlock(nn.Module): + expansion = 1 + + def __init__(self, in_planes, planes, stride=1, option='A'): + super(BasicBlock, self).__init__() + self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False) + self.bn1 = nn.BatchNorm2d(planes) + self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(planes) + + self.shortcut = nn.Sequential() + if stride != 1 or in_planes != planes: + if option == 'A': + """ + For CIFAR10 ResNet paper uses option A. + """ + self.shortcut = LambdaLayer(lambda x: + F.pad(x[:, :, ::2, ::2], (0, 0, 0, 0, planes // 4, planes // 4), "constant", + 0)) + elif option == 'B': + self.shortcut = nn.Sequential( + nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False), + nn.BatchNorm2d(self.expansion * planes) + ) + + def forward(self, x): + out = F.relu(self.bn1(self.conv1(x))) + out = self.bn2(self.conv2(out)) + out += self.shortcut(x) + out = F.relu(out) + return out + + +class ResNet(nn.Module): + def __init__(self, block, num_blocks, num_classes=10): + super(ResNet, self).__init__() + self.in_planes = 16 + + self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1, bias=False) + self.bn1 = nn.BatchNorm2d(16) + self.layer1 = self._make_layer(block, 16, num_blocks[0], stride=1) + self.layer2 = self._make_layer(block, 32, num_blocks[1], stride=2) + self.layer3 = self._make_layer(block, 64, num_blocks[2], stride=2) + self.linear = nn.Linear(64, num_classes) + + self.apply(_weights_init) + + def _make_layer(self, block, planes, num_blocks, stride): + strides = [stride] + [1] * (num_blocks - 1) + layers = [] + for stride in strides: + layers.append(block(self.in_planes, planes, stride)) + self.in_planes = planes * block.expansion + + return nn.Sequential(*layers) + + def forward(self, x): + out = F.relu(self.bn1(self.conv1(x))) + out = self.layer1(out) + out = self.layer2(out) + out = self.layer3(out) + out = F.avg_pool2d(out, out.size()[3]) + out = out.view(out.size(0), -1) + out = self.linear(out) + return out + + +def resnet20(): + return ResNet(BasicBlock, [3, 3, 3]) + + +def resnet32(): + return ResNet(BasicBlock, [5, 5, 5]) + + +def resnet44(): + return ResNet(BasicBlock, [7, 7, 7]) + + +def resnet56(): + return ResNet(BasicBlock, [9, 9, 9], num_classes=100) + + +def resnet110(): + return ResNet(BasicBlock, [18, 18, 18]) + + +def resnet1202(): + return ResNet(BasicBlock, [200, 200, 200]) + + +def test(net): + import numpy as np + total_params = 0 + + for x in filter(lambda p: p.requires_grad, net.parameters()): + total_params += np.prod(x.data.numpy().shape) + print("Total number of params", total_params) + print("Total layers", len(list(filter(lambda p: p.requires_grad and len(p.data.size()) > 1, net.parameters())))) + + +def accuracy(output, target, topk=(1,)): + """Computes the precision@k for the specified values of k""" + maxk = max(topk) + batch_size = target.size(0) + + _, pred = output.topk(maxk, 1, True, True) + pred = pred.t() + correct = pred.eq(target.view(1, -1).expand_as(pred)) + + res = [] + for k in topk: + correct_k = correct[:k].contiguous().view(-1).float().sum(0) + res.append(correct_k.mul_(100.0 / batch_size)) + return res + + +if __name__ == "__main__": + for net_name in __all__: + if net_name.startswith('resnet'): + print(net_name) + test(globals()[net_name]()) + print() diff --git a/save_checkpoint.py b/save_checkpoint.py index adb05de..341b594 100644 --- a/save_checkpoint.py +++ b/save_checkpoint.py @@ -92,10 +92,10 @@ class SaveCheckpoint(ModelCheckpoint): f.write(line + '\n') # 每次更新ckpt文件后, 将其存放到另一个位置 if self.path_final_save is not None: - zip_dir('./logs', './logs.zip') - if os.path.exists(self.path_final_save + '/logs.zip'): - os.remove(self.path_final_save + '/logs.zip') - shutil.move('./logs.zip', self.path_final_save) + zip_dir('./logs/default/' + self.dirpath.split('\\')[1], './' + self.dirpath.split('\\')[1] + '.zip') + if os.path.exists(self.path_final_save + '/' + self.dirpath.split('\\')[1] + '.zip'): + os.remove(self.path_final_save + '/' + self.dirpath.split('\\')[1] + '.zip') + shutil.move('./' + self.dirpath.split('\\')[1] + '.zip', self.path_final_save) elif self.verbose: epoch = monitor_candidates.get("epoch") step = monitor_candidates.get("step") diff --git a/train_model.py b/train_model.py index 40e2116..2e44a64 100644 --- a/train_model.py +++ b/train_model.py @@ -6,8 +6,7 @@ from pytorch_lightning.utilities.types import EPOCH_OUTPUT from torch import nn import torch -from network.MLP_JDLU import MLP_JDLU -from network.MLP_ReLU import MLP_ReLU +from network.res_net import resnet56, accuracy class TrainModule(pl.LightningModule): @@ -15,30 +14,32 @@ class TrainModule(pl.LightningModule): super().__init__() self.time_sum = None self.config = config - if 1: - self.net = MLP_ReLU(config['dim_in'], config['dim'], config['res_coef'], config['dropout_p'], - config['n_layers']) - else: - self.net = MLP_JDLU(config['dim_in'], config['dim'], config['res_coef'], config['dropout_p'], - config['n_layers']) - self.loss = nn.MSELoss() + self.net = resnet56() + self.loss = nn.CrossEntropyLoss() def training_step(self, batch, batch_idx): - x, y = batch - x = self.net(x) - loss = self.loss(x, y.type(torch.float32)) + _, input, label = batch + label = label.flatten() + pred = self.net(input) + loss = self.loss(pred, label) self.log("Training loss", loss) + acc = accuracy(pred, label)[0] + self.log("Training acc", acc) return loss def validation_step(self, batch, batch_idx): - x, y = batch - x = self.net(x) - loss = self.loss(x, y.type(torch.float32)) + _, input, label = batch + label = label.flatten() + pred = self.net(input) + loss = self.loss(pred, label) self.log("Validation loss", loss) + acc = accuracy(pred, label)[0] + self.log("Validation acc", acc) return loss def test_step(self, batch, batch_idx): - input, label = batch + _, input, label = batch + label = label.flatten() if self.time_sum is None: time_start = time.time() pred = self.net(input) @@ -47,46 +48,14 @@ class TrainModule(pl.LightningModule): print(f'\n推理时间为: {self.time_sum:f}') else: pred = self.net(input) - loss = self.loss(pred.reshape(1), label.type(torch.float32)) + loss = self.loss(pred, label) self.log("Test loss", loss) + acc = accuracy(pred, label)[0] + self.log("Test acc", acc) return input, label, pred - def test_epoch_end(self, outputs: EPOCH_OUTPUT) -> None: - records = numpy.empty((self.config['dataset_len'], 4)) - # count - for cou in range(len(outputs)): - records[cou, 0] = outputs[cou][0][0, 0] - records[cou, 1] = outputs[cou][0][0, 1] - records[cou, 2] = outputs[cou][1][0] - records[cou, 3] = outputs[cou][2] - - import plotly.graph_objects as go - trace0 = go.Mesh3d(x=records[:, 0], - y=records[:, 1], - z=records[:, 2], - opacity=0.5, - name='label' - ) - trace1 = go.Mesh3d(x=records[:, 0], - y=records[:, 1], - z=records[:, 3], - opacity=0.5, - name='pred' - ) - fig = go.Figure(data=[trace0, trace1]) - fig.update_layout( - scene=dict( - # xaxis=dict(nticks=4, range=[-100, 100], ), - # yaxis=dict(nticks=4, range=[-50, 100], ), - # zaxis=dict(nticks=4, range=[-100, 100], ), - aspectratio=dict(x=1, y=1, z=0.5), - ), - - ) - fig.show() - def configure_optimizers(self): - optimizer = torch.optim.Adam(self.parameters(), lr=1e-3) + optimizer = torch.optim.SGD(self.parameters(), lr=0.1, momentum=0.9, weight_decay=1e-4) return optimizer def load_pretrain_parameters(self):