From 3605b974abd8a05357a4ac4bd7f981b0c65e7cfd Mon Sep 17 00:00:00 2001 From: flt6 <1404262047@qq.com> Date: Fri, 16 May 2025 15:48:59 +0800 Subject: [PATCH] add source single test for train Former-commit-id: bbbc19eaa23049a6edabf93731f799f391e9623a --- Picture_Train/model.py | 198 ++++++++++++++++++++++++++++++++ Picture_Train/predict.py | 6 +- Picture_Train/resnet34-1Net.pth | 3 + Picture_Train/src_predict.py | 40 +++++++ 4 files changed, 244 insertions(+), 3 deletions(-) create mode 100644 Picture_Train/model.py create mode 100644 Picture_Train/resnet34-1Net.pth create mode 100644 Picture_Train/src_predict.py diff --git a/Picture_Train/model.py b/Picture_Train/model.py new file mode 100644 index 0000000..c2aa086 --- /dev/null +++ b/Picture_Train/model.py @@ -0,0 +1,198 @@ +import torch.nn as nn +import torch + + +class BasicBlock(nn.Module): + expansion = 1 + + def __init__(self, in_channel, out_channel, stride=1, downsample=None, **kwargs): + super(BasicBlock, self).__init__() + self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=out_channel, + kernel_size=3, stride=stride, padding=1, bias=False) + self.bn1 = nn.BatchNorm2d(out_channel) + self.relu = nn.ReLU() + self.conv2 = nn.Conv2d(in_channels=out_channel, out_channels=out_channel, + kernel_size=3, stride=1, padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(out_channel) + self.downsample = downsample + + def forward(self, x): + identity = x + if self.downsample is not None: + identity = self.downsample(x) + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + + out += identity + out = self.relu(out) + + return out + + +class Bottleneck(nn.Module): + """ + 注意：原论文中，在虚线残差结构的主分支上，第一个1x1卷积层的步距是2，第二个3x3卷积层步距是1。 + 但在pytorch官方实现过程中是第一个1x1卷积层的步距是1，第二个3x3卷积层步距是2， + 这么做的好处是能够在top1上提升大概0.5%的准确率。 + 可参考Resnet v1.5 https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch + """ + expansion = 4 + + def __init__(self, in_channel, out_channel, stride=1, downsample=None, + groups=1, width_per_group=64): + super(Bottleneck, self).__init__() + + width = int(out_channel * (width_per_group / 64.)) * groups + + self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=width, + kernel_size=1, stride=1, bias=False) # squeeze channels + self.bn1 = nn.BatchNorm2d(width) + # ----------------------------------------- + self.conv2 = nn.Conv2d(in_channels=width, out_channels=width, groups=groups, + kernel_size=3, stride=stride, bias=False, padding=1) + self.bn2 = nn.BatchNorm2d(width) + # ----------------------------------------- + self.conv3 = nn.Conv2d(in_channels=width, out_channels=out_channel*self.expansion, + kernel_size=1, stride=1, bias=False) # unsqueeze channels + self.bn3 = nn.BatchNorm2d(out_channel*self.expansion) + self.relu = nn.ReLU(inplace=True) + self.downsample = downsample + + def forward(self, x): + identity = x + if self.downsample is not None: + identity = self.downsample(x) + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.relu(out) + + out = self.conv3(out) + out = self.bn3(out) + + out += identity + out = self.relu(out) + + return out + + +class ResNet(nn.Module): + + def __init__(self, + block, + blocks_num, + num_classes=1000, + include_top=True, + groups=1, + width_per_group=64): + super(ResNet, self).__init__() + self.include_top = include_top + self.in_channel = 64 + + self.groups = groups + self.width_per_group = width_per_group + + self.conv1 = nn.Conv2d(3, self.in_channel, kernel_size=7, stride=2, + padding=3, bias=False) + self.bn1 = nn.BatchNorm2d(self.in_channel) + self.relu = nn.ReLU(inplace=True) + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + self.layer1 = self._make_layer(block, 64, blocks_num[0]) + self.layer2 = self._make_layer(block, 128, blocks_num[1], stride=2) + self.layer3 = self._make_layer(block, 256, blocks_num[2], stride=2) + self.layer4 = self._make_layer(block, 512, blocks_num[3], stride=2) + if self.include_top: + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) # output size = (1, 1) + self.fc = nn.Linear(512 * block.expansion, num_classes) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + + def _make_layer(self, block, channel, block_num, stride=1): + downsample = None + if stride != 1 or self.in_channel != channel * block.expansion: + downsample = nn.Sequential( + nn.Conv2d(self.in_channel, channel * block.expansion, kernel_size=1, stride=stride, bias=False), + nn.BatchNorm2d(channel * block.expansion)) + + layers = [] + layers.append(block(self.in_channel, + channel, + downsample=downsample, + stride=stride, + groups=self.groups, + width_per_group=self.width_per_group)) + self.in_channel = channel * block.expansion + + for _ in range(1, block_num): + layers.append(block(self.in_channel, + channel, + groups=self.groups, + width_per_group=self.width_per_group)) + + return nn.Sequential(*layers) + + def forward(self, x): + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + x = self.maxpool(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.layer4(x) + + if self.include_top: + x = self.avgpool(x) + x = torch.flatten(x, 1) + x = self.fc(x) + + return x + + +def resnet34(num_classes=1000, include_top=True): + # https://download.pytorch.org/models/resnet34-333f7ec4.pth + return ResNet(BasicBlock, [3, 4, 6, 3], num_classes=num_classes, include_top=include_top) + + +def resnet50(num_classes=1000, include_top=True): + # https://download.pytorch.org/models/resnet50-19c8e357.pth + return ResNet(Bottleneck, [3, 4, 6, 3], num_classes=num_classes, include_top=include_top) + + +def resnet101(num_classes=1000, include_top=True): + # https://download.pytorch.org/models/resnet101-5d3b4d8f.pth + return ResNet(Bottleneck, [3, 4, 23, 3], num_classes=num_classes, include_top=include_top) + + +def resnext50_32x4d(num_classes=1000, include_top=True): + # https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth + groups = 32 + width_per_group = 4 + return ResNet(Bottleneck, [3, 4, 6, 3], + num_classes=num_classes, + include_top=include_top, + groups=groups, + width_per_group=width_per_group) + + +def resnext101_32x8d(num_classes=1000, include_top=True): + # https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth + groups = 32 + width_per_group = 8 + return ResNet(Bottleneck, [3, 4, 23, 3], + num_classes=num_classes, + include_top=include_top, + groups=groups, + width_per_group=width_per_group) diff --git a/Picture_Train/predict.py b/Picture_Train/predict.py index 101adce..4050bc7 100644 --- a/Picture_Train/predict.py +++ b/Picture_Train/predict.py @@ -6,7 +6,7 @@ from PIL import Image from torchvision import transforms import matplotlib.pyplot as plt -from model import vgg +from model_train import vgg # 单张图片预测程序 def main(): @@ -18,7 +18,7 @@ def main(): transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) # load image - img_path = "C:/Picture_Train/test-data/pic_0032.jpg" + img_path = "data/val/yellow/ticket_0_3.jpg" assert os.path.exists(img_path), "file: '{}' dose not exist.".format(img_path) img = Image.open(img_path) plt.imshow(img) @@ -37,7 +37,7 @@ def main(): # create model model = vgg(model_name="vgg16", num_classes=2).to(device) # load model weights - weights_path = "C:/Picture_Train/resnet34Net.pth" + weights_path = "resnet34-1Net.pth" assert os.path.exists(weights_path), "file: '{}' dose not exist.".format(weights_path) model.load_state_dict(torch.load(weights_path, map_location=device)) diff --git a/Picture_Train/resnet34-1Net.pth b/Picture_Train/resnet34-1Net.pth new file mode 100644 index 0000000..a42d612 --- /dev/null +++ b/Picture_Train/resnet34-1Net.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:eded1e9cf905cceb40b6eadf6ec9979852f4ddaf2d1d470cf681f1fbf0d3a03a +size 85279422 diff --git a/Picture_Train/src_predict.py b/Picture_Train/src_predict.py new file mode 100644 index 0000000..fea8c4a --- /dev/null +++ b/Picture_Train/src_predict.py @@ -0,0 +1,40 @@ +from PIL import Image +from torchvision import transforms +import torch +import os +import json +from model import resnet34 + +print("init model") +model = resnet34(num_classes=2).to("cpu") + +assert os.path.exists("resnet34-1Net.pth"), "file: '{}' dose not exist.".format("resnet34-1Net.pth") +model.load_state_dict(torch.load("resnet34-1Net.pth", map_location="cpu")) + +print("load model done") + +def predictor(im_file): # 预测分类 + image = Image.open(im_file) + data_transform = transforms.Compose([ + transforms.Resize(256), + transforms.CenterCrop(224), + transforms.ToTensor(), + transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + ]) + img = data_transform(image) + img = torch.unsqueeze(img, dim=0) + with open("class_indices.json", "r") as f: + class_indict = json.load(f) + + model.eval() + with torch.no_grad(): + output = torch.squeeze(model(img)).cpu() + predict = torch.softmax(output, dim=0) + predict_cla = torch.argmax(predict).numpy() + + class_a = "{}".format(class_indict[str(predict_cla)]) + prob_a = "{:.3}".format(predict[predict_cla].numpy()) + prob_b = float(prob_a) + # print('class_:',class_a) + # print('prob_:',prob_b) + return class_a, prob_b \ No newline at end of file