Pytorch图像处理篇：使用pytorch搭建ResNet并基于迁移学习训练

model.py

import torch.nn as nn
import torch#首先定义34层残差结构
class BasicBlock(nn.Module):expansion &＃61; 1 #对应主分支中卷积核的个数有没有发生变化#定义初始化函数&＃xff08;输入特征矩阵的深度&＃xff0c;输出特征矩阵的深度&＃xff08;主分支上卷积核的个数&＃xff09;&＃xff0c;不惧默认设置为1&＃xff0c;下采样参数设置为None&＃xff09;def __init__(self, in_channel, out_channel, stride&＃61;1, downsample&＃61;None, **kwargs):super(BasicBlock, self).__init__()self.conv1 &＃61; nn.Conv2d(in_channels&＃61;in_channel, out_channels&＃61;out_channel,kernel_size&＃61;3, stride&＃61;stride, padding&＃61;1, bias&＃61;False)self.bn1 &＃61; nn.BatchNorm2d(out_channel)self.relu &＃61; nn.ReLU()self.conv2 &＃61; nn.Conv2d(in_channels&＃61;out_channel, out_channels&＃61;out_channel,kernel_size&＃61;3, stride&＃61;1, padding&＃61;1, bias&＃61;False)self.bn2 &＃61; nn.BatchNorm2d(out_channel)self.downsample &＃61; downsample#定义正向传播的过程def forward(self, x):identity &＃61; xif self.downsample is not None:identity &＃61; self.downsample(x) #将输入传入下采样函数得到捷径分支的输出#主分支上的输出out &＃61; self.conv1(x)out &＃61; self.bn1(out)out &＃61; self.relu(out)out &＃61; self.conv2(out)out &＃61; self.bn2(out)out &＃43;&＃61; identity #将主分支上的输出加上捷径分支上的输出out &＃61; self.relu(out)return out #得到残差结构的最终输出#定义50层、101层、152层的残差结构&＃xff0c;在这个网络上进行修改得到ResNext网络
class Bottleneck(nn.Module):"""注意&＃xff1a;原论文中&＃xff0c;在虚线残差结构的主分支上&＃xff0c;第一个1x1卷积层的步距是2&＃xff0c;第二个3x3卷积层步距是1。但在pytorch官方实现过程中是第一个1x1卷积层的步距是1&＃xff0c;第二个3x3卷积层步距是2&＃xff0c;这么做的好处是能够在top1上提升大概0.5%的准确率。可参考Resnet v1.5 https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch"""expansion &＃61; 4 #残差结构所使用卷积核的一个变化#定义初始化函数def __init__(self, in_channel, out_channel, stride&＃61;1, downsample&＃61;None,#相比resnet网络多传入了两个参数groups&＃61;1, width_per_group&＃61;64&＃xff0c;groups&＃61;1, width_per_group&＃61;64):super(Bottleneck, self).__init__()width &＃61; int(out_channel * (width_per_group / 64.)) * groupsself.conv1 &＃61; nn.Conv2d(in_channels&＃61;in_channel, out_channels&＃61;width,kernel_size&＃61;1, stride&＃61;1, bias&＃61;False) # squeeze channelsself.bn1 &＃61; nn.BatchNorm2d(width)# -----------------------------------------#输入、输出特征矩阵的channel设置为widthself.conv2 &＃61; nn.Conv2d(in_channels&＃61;width, out_channels&＃61;width, groups&＃61;groups,kernel_size&＃61;3, stride&＃61;stride, bias&＃61;False, padding&＃61;1)self.bn2 &＃61; nn.BatchNorm2d(width)# -----------------------------------------self.conv3 &＃61; nn.Conv2d(in_channels&＃61;width, out_channels&＃61;out_channel*self.expansion,kernel_size&＃61;1, stride&＃61;1, bias&＃61;False) # unsqueeze channelsself.bn3 &＃61; nn.BatchNorm2d(out_channel*self.expansion)self.relu &＃61; nn.ReLU(inplace&＃61;True)self.downsample &＃61; downsample#定义正向传播过程def forward(self, x):identity &＃61; xif self.downsample is not None:identity &＃61; self.downsample(x)out &＃61; self.conv1(x)out &＃61; self.bn1(out)out &＃61; self.relu(out)out &＃61; self.conv2(out)out &＃61; self.bn2(out)out &＃61; self.relu(out)out &＃61; self.conv3(out)out &＃61; self.bn3(out)out &＃43;&＃61; identityout &＃61; self.relu(out)return out#定义ResNet网络模型
class ResNet(nn.Module):def __init__(self,block,#对应的就是残差结构blocks_num, #所使用残差结构的数目num_classes&＃61;1000, #训练集的分类个数include_top&＃61;True,#是为了在ResNet网络上搭建更复杂的网络groups&＃61;1,width_per_group&＃61;64):super(ResNet, self).__init__()self.include_top &＃61; include_top #传入类变量之中self.in_channel &＃61; 64 #输入特征矩阵的深度self.groups &＃61; groupsself.width_per_group &＃61; width_per_group#定义第一层的卷积层&＃xff0c;3表示输入矩阵的深度self.conv1 &＃61; nn.Conv2d(3, self.in_channel, kernel_size&＃61;7, stride&＃61;2,padding&＃61;3, bias&＃61;False)self.bn1 &＃61; nn.BatchNorm2d(self.in_channel)self.relu &＃61; nn.ReLU(inplace&＃61;True)self.maxpool &＃61; nn.MaxPool2d(kernel_size&＃61;3, stride&＃61;2, padding&＃61;1) #最大池化下采样操作self.layer1 &＃61; self._make_layer(block, 64, blocks_num[0]) #一系列残差结构self.layer2 &＃61; self._make_layer(block, 128, blocks_num[1], stride&＃61;2)self.layer3 &＃61; self._make_layer(block, 256, blocks_num[2], stride&＃61;2)self.layer4 &＃61; self._make_layer(block, 512, blocks_num[3], stride&＃61;2)if self.include_top:self.avgpool &＃61; nn.AdaptiveAvgPool2d((1, 1)) # 自适应的平均池化下采样output size &＃61; (1, 1)self.fc &＃61; nn.Linear(512 * block.expansion, num_classes) #全连接层for m in self.modules():if isinstance(m, nn.Conv2d):nn.init.kaiming_normal_(m.weight, mode&＃61;&＃39;fan_out&＃39;, nonlinearity&＃61;&＃39;relu&＃39;)#&＃xff08;哪一个残差结构&＃xff0c;残差结构中第一卷积层所使用卷积核的个数&＃xff0c;该层包含了几个残差结构&＃xff0c;步距为1&＃xff09;def _make_layer(self, block, channel, block_num, stride&＃61;1):downsample &＃61; None #定义下采样if stride !&＃61; 1 or self.in_channel !&＃61; channel * block.expansion: #对于十八层和三十四层的残差结构&＃xff0c;就会跳过if语句&＃xff1b;downsample &＃61; nn.Sequential( #生成下采样函数nn.Conv2d(self.in_channel, channel * block.expansion, kernel_size&＃61;1, stride&＃61;stride, bias&＃61;False),nn.BatchNorm2d(channel * block.expansion))layers &＃61; [] #定义空的列表#将第一层的残差结构传进去layers.append(block(self.in_channel,channel,downsample&＃61;downsample,stride&＃61;stride,groups&＃61;self.groups,width_per_group&＃61;self.width_per_group))self.in_channel &＃61; channel * block.expansion#实现实线部分for _ in range(1, block_num): #表示从一开始遍历&＃xff0c;不写则默认是0层开始layers.append(block(self.in_channel,channel,groups&＃61;self.groups,width_per_group&＃61;self.width_per_group))return nn.Sequential(*layers) #非关键字参数的方式传入nn.squential函数#进行正向传播过程def forward(self, x):x &＃61; self.conv1(x)x &＃61; self.bn1(x)x &＃61; self.relu(x)x &＃61; self.maxpool(x)x &＃61; self.layer1(x)x &＃61; self.layer2(x)x &＃61; self.layer3(x)x &＃61; self.layer4(x)if self.include_top:x &＃61; self.avgpool(x)x &＃61; torch.flatten(x, 1)x &＃61; self.fc(x)return x#对网络进行实例化&＃xff0c;传入BasicBlock或者Bottleneck来确定是哪个网络&＃xff0c;第二个参数是block的个数
def resnet34(num_classes&＃61;1000, include_top&＃61;True):# https://download.pytorch.org/models/resnet34-333f7ec4.pthreturn ResNet(BasicBlock, [3, 4, 6, 3], num_classes&＃61;num_classes, include_top&＃61;include_top)def resnet50(num_classes&＃61;1000, include_top&＃61;True):# https://download.pytorch.org/models/resnet50-19c8e357.pthreturn ResNet(Bottleneck, [3, 4, 6, 3], num_classes&＃61;num_classes, include_top&＃61;include_top)def resnet101(num_classes&＃61;1000, include_top&＃61;True):# https://download.pytorch.org/models/resnet101-5d3b4d8f.pthreturn ResNet(Bottleneck, [3, 4, 23, 3], num_classes&＃61;num_classes, include_top&＃61;include_top)#进行升级ResNext网络
def resnext50_32x4d(num_classes&＃61;1000, include_top&＃61;True):# https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pthgroups &＃61; 32width_per_group &＃61; 4return ResNet(Bottleneck, [3, 4, 6, 3],num_classes&＃61;num_classes,include_top&＃61;include_top,groups&＃61;groups,width_per_group&＃61;width_per_group)def resnext101_32x8d(num_classes&＃61;1000, include_top&＃61;True):# https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pthgroups &＃61; 32width_per_group &＃61; 8return ResNet(Bottleneck, [3, 4, 23, 3],num_classes&＃61;num_classes,include_top&＃61;include_top,groups&＃61;groups,width_per_group&＃61;width_per_group)


train.py

import os
import sys
import jsonimport torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms, datasets
from tqdm import tqdm
from model import resnet34 #要进行训练的话&＃xff0c;要导入需要的网络&＃xff0c;是resnet34还是rtesnet50或者其他网络def main():device &＃61; torch.device("cuda:0" if torch.cuda.is_available() else "cpu")print("using {} device.".format(device))data_transform &＃61; {"train": transforms.Compose([transforms.RandomResizedCrop(224),transforms.RandomHorizontalFlip(),transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),#进行标准化的方法的参数是参考官网"val": transforms.Compose([transforms.Resize(256),#先通过resize将最小遍缩放到256&＃xff0c;transforms.CenterCrop(224),#在使用中心裁剪transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])}data_root &＃61; os.path.abspath(os.path.join(os.getcwd(), "../..")) # get data root pathimage_path &＃61; os.path.join(data_root, "data_set", "flower_data") # flower data set pathassert os.path.exists(image_path), "{} path does not exist.".format(image_path)train_dataset &＃61; datasets.ImageFolder(root&＃61;os.path.join(image_path, "train"),transform&＃61;data_transform["train"])train_num &＃61; len(train_dataset)# {&＃39;daisy&＃39;:0, &＃39;dandelion&＃39;:1, &＃39;roses&＃39;:2, &＃39;sunflower&＃39;:3, &＃39;tulips&＃39;:4}flower_list &＃61; train_dataset.class_to_idxcla_dict &＃61; dict((val, key) for key, val in flower_list.items())# write dict into json filejson_str &＃61; json.dumps(cla_dict, indent&＃61;4)with open(&＃39;class_indices.json&＃39;, &＃39;w&＃39;) as json_file:json_file.write(json_str)batch_size &＃61; 16nw &＃61; min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8]) # number of workersprint(&＃39;Using {} dataloader workers every process&＃39;.format(nw))train_loader &＃61; torch.utils.data.DataLoader(train_dataset,batch_size&＃61;batch_size, shuffle&＃61;True,num_workers&＃61;nw)validate_dataset &＃61; datasets.ImageFolder(root&＃61;os.path.join(image_path, "val"),transform&＃61;data_transform["val"])val_num &＃61; len(validate_dataset)validate_loader &＃61; torch.utils.data.DataLoader(validate_dataset,batch_size&＃61;batch_size, shuffle&＃61;False,num_workers&＃61;nw)print("using {} images for training, {} images for validation.".format(train_num,val_num))#使用迁移学习的方法&＃xff0c;所以需要使用Pytorch官方所提供的resnet网络的预训练模型&＃xff0c;需要去下载net &＃61; resnet34() #需要用哪个网络就实例化哪个网络# load pretrain weights# download url: https://download.pytorch.org/models/resnet34-333f7ec4.pthmodel_weight_path &＃61; "./resnet34-pre.pth" #使用哪个权重就写哪个assert os.path.exists(model_weight_path), "file {} does not exist.".format(model_weight_path)net.load_state_dict(torch.load(model_weight_path, map_location&＃61;&＃39;cpu&＃39;))# for param in net.parameters():# param.requires_grad &＃61; False#载入预训练模型的方法# change fc layer structurein_channel &＃61; net.fc.in_featuresnet.fc &＃61; nn.Linear(in_channel, 5) #将最后一个新连接层&＃xff0c;替换成自己的新建的一个全连接层&＃xff0c;5表示要分类的类别个数net.to(device)# define loss functionloss_function &＃61; nn.CrossEntropyLoss()# construct an optimizerparams &＃61; [p for p in net.parameters() if p.requires_grad]optimizer &＃61; optim.Adam(params, lr&＃61;0.0001)epochs &＃61; 3best_acc &＃61; 0.0save_path &＃61; &＃39;./resNet34.pth&＃39; #保存权重的名字也进行相应的修改train_steps &＃61; len(train_loader)for epoch in range(epochs):# trainnet.train()#重要的running_loss &＃61; 0.0train_bar &＃61; tqdm(train_loader, file&＃61;sys.stdout)for step, data in enumerate(train_bar):images, labels &＃61; dataoptimizer.zero_grad()logits &＃61; net(images.to(device))loss &＃61; loss_function(logits, labels.to(device))loss.backward()optimizer.step()# print statisticsrunning_loss &＃43;&＃61; loss.item()train_bar.desc &＃61; "train epoch[{}/{}] loss:{:.3f}".format(epoch &＃43; 1,epochs,loss)# validatenet.eval()acc &＃61; 0.0 # accumulate accurate number / epochwith torch.no_grad():val_bar &＃61; tqdm(validate_loader, file&＃61;sys.stdout)for val_data in val_bar:val_images, val_labels &＃61; val_dataoutputs &＃61; net(val_images.to(device))# loss &＃61; loss_function(outputs, test_labels)predict_y &＃61; torch.max(outputs, dim&＃61;1)[1]acc &＃43;&＃61; torch.eq(predict_y, val_labels.to(device)).sum().item()val_bar.desc &＃61; "valid epoch[{}/{}]".format(epoch &＃43; 1,epochs)val_accurate &＃61; acc / val_numprint(&＃39;[epoch %d] train_loss: %.3f val_accuracy: %.3f&＃39; %(epoch &＃43; 1, running_loss / train_steps, val_accurate))if val_accurate > best_acc:best_acc &＃61; val_accuratetorch.save(net.state_dict(), save_path)print(&＃39;Finished Training&＃39;)if __name__ &＃61;&＃61; &＃39;__main__&＃39;:main()


predict.py

import os
import jsonimport torch
from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as pltfrom model import resnet34def main():device &＃61; torch.device("cuda:0" if torch.cuda.is_available() else "cpu")data_transform &＃61; transforms.Compose([transforms.Resize(256),transforms.CenterCrop(224),transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])# load imageimg_path &＃61; "../tulip.jpg"assert os.path.exists(img_path), "file: &＃39;{}&＃39; dose not exist.".format(img_path)img &＃61; Image.open(img_path)plt.imshow(img)# [N, C, H, W]img &＃61; data_transform(img)# expand batch dimensionimg &＃61; torch.unsqueeze(img, dim&＃61;0)# read class_indictjson_path &＃61; &＃39;./class_indices.json&＃39;assert os.path.exists(json_path), "file: &＃39;{}&＃39; dose not exist.".format(json_path)with open(json_path, "r") as f:class_indict &＃61; json.load(f)# create model 使用哪个网络就传入哪个网络model &＃61; resnet34(num_classes&＃61;5).to(device)# load model weightsweights_path &＃61; "./resNet34.pth" #权重进行相应的改变assert os.path.exists(weights_path), "file: &＃39;{}&＃39; dose not exist.".format(weights_path)model.load_state_dict(torch.load(weights_path, map_location&＃61;device))# predictionmodel.eval()with torch.no_grad():#不对损失梯度进行跟踪# predict classoutput &＃61; torch.squeeze(model(img.to(device))).cpu()predict &＃61; torch.softmax(output, dim&＃61;0)predict_cla &＃61; torch.argmax(predict).numpy() #最大值对相应的索引print_res &＃61; "class: {} prob: {:.3}".format(class_indict[str(predict_cla)],predict[predict_cla].numpy())plt.title(print_res)for i in range(len(predict)):print("class: {:10} prob: {:.3}".format(class_indict[str(i)],predict[i].numpy()))plt.show()if __name__ &＃61;&＃61; &＃39;__main__&＃39;:main()


batch_predict.py


#批量进行预测
import os
import jsonimport torch
from PIL import Image
from torchvision import transformsfrom model import resnet34def main():device &＃61; torch.device("cuda:0" if torch.cuda.is_available() else "cpu")data_transform &＃61; transforms.Compose([transforms.Resize(256),transforms.CenterCrop(224),transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])# load image# 指向需要遍历预测的图像文件夹imgs_root &＃61; "/data/imgs"assert os.path.exists(imgs_root), f"file: &＃39;{imgs_root}&＃39; dose not exist."# 读取指定文件夹下所有jpg图像路径img_path_list &＃61; [os.path.join(imgs_root, i) for i in os.listdir(imgs_root) if i.endswith(".jpg")]# read class_indictjson_path &＃61; &＃39;./class_indices.json&＃39;assert os.path.exists(json_path), f"file: &＃39;{json_path}&＃39; dose not exist."json_file &＃61; open(json_path, "r")class_indict &＃61; json.load(json_file)# create modelmodel &＃61; resnet34(num_classes&＃61;5).to(device)# load model weightsweights_path &＃61; "./resNet34.pth"assert os.path.exists(weights_path), f"file: &＃39;{weights_path}&＃39; dose not exist."model.load_state_dict(torch.load(weights_path, map_location&＃61;device))# predictionmodel.eval()batch_size &＃61; 8 # 每次预测时将多少张图片打包成一个batchwith torch.no_grad():for ids in range(0, len(img_path_list) // batch_size):img_list &＃61; []for img_path in img_path_list[ids * batch_size: (ids &＃43; 1) * batch_size]:assert os.path.exists(img_path), f"file: &＃39;{img_path}&＃39; dose not exist."img &＃61; Image.open(img_path)img &＃61; data_transform(img)img_list.append(img)# batch img# 将img_list列表中的所有图像打包成一个batchbatch_img &＃61; torch.stack(img_list, dim&＃61;0)# predict classoutput &＃61; model(batch_img.to(device)).cpu()predict &＃61; torch.softmax(output, dim&＃61;1)probs, classes &＃61; torch.max(predict, dim&＃61;1)for idx, (pro, cla) in enumerate(zip(probs, classes)):print("image: {} class: {} prob: {:.3}".format(img_path_list[ids * batch_size &＃43; idx],class_indict[str(cla.numpy())],pro.numpy()))if __name__ &＃61;&＃61; &＃39;__main__&＃39;:main()