赞
踩
想入门pytorch强化学习,就去找pytorch的课来看。B站上播放量最高的就是小土堆的课,整体跟下来感觉内容还是很详细的,但和我的预期不太一样,这个是DL的不是RL的,不过作为对于pytorch使用的初期了解也是很好的,这篇博客就把整个学习过程做一个梳理。
注意:本笔记使用的数据集全部都是CIFAR10,下载比较简单~,下面开始
在读取之前,需要先准备好数据了,对于CIFAR10,可以离线下载(网址:https://download.pytorch.org/tutorial/hymenoptera_data.zip),下载后保存到dataset文件夹,目录结构如下:

下面就是对于图片的读取,起初的读取都是通过PIL,后面会换成dataloader,主要是自己定义了一个类,传递了两个参数,实例如下:
# function:使用PIL完成数据的读取,可查看 from torch.utils.data import Dataset from PIL import Image import os class MyData(Dataset): def __init__(self, root_dir, label_dir): self.root_dir = root_dir self.label_dir = label_dir self.path = os.path.join(self.root_dir, self.label_dir) self.img_path = os.listdir(self.path) def __getitem__(self, idx): img_name = self.img_path[idx] img_item_path = os.path.join(self.root_dir, self.label_dir, img_name) img = Image.open(img_item_path) img.show() label = self.label_dir return img, label def __len__(self): return len(self.img_path) root_dir = "dataset/train" ants_label_dir = "ants" ants_datasets = MyData(root_dir, ants_label_dir) ants_datasets.__getitem__(0) # 输入查看图片编号即可 print(len(ants_datasets))
tensorboard是一个可视化工具,可以用来看图片或者分析数据。
先说一下安装,在安装的时候报了两个错:
报错:ModuleNotFoundError: No module named 'tensorboard'
解决:pip install tensorboard -i http://mirrors.aliyun.com/pypi/simple --trusted-host mirrors.aliyun.com
报错:ModuleNotFoundError: No module named 'six'
解决:pip install six -i http://mirrors.aliyun.com/pypi/simple --trusted-host mirrors.aliyun.com
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter("logs") # 这个logs是生成日志文件的文件夹名,可随意更换
writer.add_image() # 添加一张图像,第一个参数是tag,第二个是数据本身,第三个是编号
wirter.add_images() # 添加多张图像,第一个是tag,第二个是数据本身,第三个是编号
writer.add_scalar() # 添加数值 第一个参数是tag, 再后是先y后x, 例如:writer.add_scalar("y=2x", 2*i, i)
writer.add_graph(tudui, input) # 查看网络结构,tudui是模型,input是模型输入
writer.close()
执行代码过后就会在logs文件夹下生成event文件:



还可以看网络结构:

# function:展示tensorboard的使用 from torch.utils.tensorboard import SummaryWriter import numpy as np from PIL import Image writer = SummaryWriter("logs") # add_image使用 image_path = "dataset/train/ants/5650366_e22b7e1065.jpg" image_path1 = "dataset/train/ants/6240329_72c01e663e.jpg" img_PIL = Image.open(image_path) img_array = np.array(img_PIL) img_PIL1 = Image.open(image_path1) img_array1 = np.array(img_PIL1) writer.add_image("image", img_array, 1, dataformats='HWC') # 这个通道顺序需要改,tensorboard默认使用的是tensor结构,但这个读的是PIL结构 writer.add_image("image", img_array1, 2, dataformats='HWC') # add_scalar使用 # for i in range(100): # writer.add_scalar("y=2x", 2*i, i) writer.close()
transforms是用来进行数据类型的转换,pytorch中使用的数据格式大多是Tensor,transfroms直接提供了工具
# function:展示transforms的基本使用格式 from PIL import Image from torch.utils.tensorboard import SummaryWriter from torchvision import transforms import numpy as np writer = SummaryWriter("logs") # python用法 → tensor数据类型 # 通过transform.ToTensor看两个问题 # 1. transform如何使用 # 2. 为什么需要Tensor image_path = "dataset/train/ants/522163566_fec115ca66.jpg" # PIL → Tensor() img = Image.open(image_path) tensor_trans = transforms.ToTensor() tensor_img = tensor_trans(img) writer.add_image("image", tensor_img, 1) writer.close()
有几个常见的transforms,可以记一下使用方法:ToTensor(转换为Tensor类型),Normalize(做正则化),Resize(调整数据shape),Compose(将多个transfroms整合在一起),RandomCrop(随机裁剪数据),使用方式如下:
# function:展示部分常见transfroms,包括:ToTensor(转换为Tensor类型),Normalize(做正则化),Resize(调整数据shape),Compose(将多个transfroms整合在一起),RandomCrop(随机裁剪数据) from PIL import Image from torch.utils.tensorboard import SummaryWriter from torchvision import transforms writer = SummaryWriter("logs") img = Image.open("dataset/train/ants/5650366_e22b7e1065.jpg") # ToTensor trans_totensor = transforms.ToTensor() img_tensor = trans_totensor(img) # Normalize trans_norm = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) img_norm = trans_norm(img_tensor) # Resize # 输入序列(512, 512)或者数值(512, 会生成方阵) trans_resize = transforms.Resize((512, 512)) img_resize = trans_resize(img) img_resize = trans_totensor(img_resize) # Compose # PIL → PIL → tensor trans_resize_2 = transforms.Resize(512) trans_compose = transforms.Compose([trans_resize_2, trans_totensor]) img_resize_2 = trans_compose(img) # RandomCrop——随机裁剪 trans_ramdom = transforms.RandomCrop(128) trans_compose_2 = transforms.Compose([trans_ramdom, trans_totensor]) for i in range(10): img_crop = trans_compose_2(img) writer.add_image("RandomCrop", img_crop, i) writer.add_image("ToTensor", img_tensor, 1) writer.add_image("Norm", img_norm, 2) writer.add_image("Resize", img_resize, 3) writer.add_image("Compose", img_resize_2, 4) writer.close()
这个主要开始使用dataLoader进行数据提取了,因为图片数据直接是PIL类型,所以在dataloader的时候就要进行transforms。datasets来源于torchversion:
torchversion:可以下载默认数据集
torchvision:dataset 下载数据集
torchvision:dataloader 选择特定数据下载
使用方式如下:
import torchvision from torch.utils.tensorboard import SummaryWriter from torch.utils.data import DataLoader # # ====torchvision.datasets使用==== # dataset_transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor()]) # train_set = torchvision.datasets.CIFAR10(root="./dataset", train=True, transform=dataset_transform, download=True) # test_set = torchvision.datasets.CIFAR10(root="./dataset", train=False, transform=dataset_transform, download=True) # # writer = SummaryWriter("dataset_transformer") # for i in range(10): # img, target = train_set[i] # writer.add_image("test_set", img, i) # writer.close() # # ====torchvision.DataLoader使用==== test_data = torchvision.datasets.CIFAR10("./dataset", train=False, transform=torchvision.transforms.Compose([torchvision.transforms.ToTensor()])) test_loader = DataLoader(dataset=test_data, batch_size=64, shuffle=True, num_workers=0, drop_last=True) img, target = test_data[0] print(img.shape) print(target) step = 0 writer = SummaryWriter("dataloader") for data in test_loader: imgs, targets = data writer.add_images("dataloader", imgs, step) # 注意:批量添加的时候使用add_images函数 step = step + 1 writer.close()
torch.nn是pytorch对于神经网络(neural network)提供的有关操作支持,具体的东西有很多,只讲解了一部分常用的
卷积操作的实现(具体啥是卷积,视频里解释的很详细),我在这里就留个模板了:
# torch.nn 的卷积conv2d 实例 import torch import torchvision from torch import nn from torch.nn import Conv2d from torch.utils.data import DataLoader from torch.utils.tensorboard import SummaryWriter dataset = torchvision.datasets.CIFAR10("./dataset", train=False, transform=torchvision.transforms.ToTensor(), download=True) dataLoader = DataLoader(dataset, batch_size=64) class Tudui(nn.Module): def __init__(self): super(Tudui, self).__init__() self.conv1 = Conv2d(in_channels=3, out_channels=6, kernel_size=3, stride=1, padding=0) def forward(self, x): x = self.conv1(x) return x tudui = Tudui() step = 0 writter = SummaryWriter("logs_conv2d") for data in dataLoader: imgs, targets = data output = tudui(imgs) writter.add_images("input", imgs, step) output = torch.reshape(output, [-1, 3, 30, 30]) writter.add_images("output", output, step) step = step + 1
上面的Tudui类继承了nn.module,其实这个是整体pytorch关于神经网络的父类,使用时继承就好了,留一个简单的模板:
# function:nn.module的基本使用 import torch from torch import nn class Tudui(nn.Module): # 注意父类写的格式 def __init__(self): super().__init__() def forward(self, input): output = input + 1 return output tudui = Tudui() x = torch.tensor(1.0) output = tudui(x) print(output)
最大池化和卷积区别
卷积是利用卷积核做计算,维度不变
卷积是利用卷积核计算后取最大值,维度不变
最大池化的实现:
# function:torch.nn 最大池化maxpool示例 import torch import torchvision from torch import nn from torch.nn import MaxPool2d from torch.utils.data import DataLoader from torch.utils.tensorboard import SummaryWriter dataset = torchvision.datasets.CIFAR10("./dataset", download=True, train=False, transform=torchvision.transforms.ToTensor()) dataLoader = DataLoader(dataset, batch_size=64) input = torch.tensor([[1, 2, 0, 3, 1], [0, 1, 2, 3, 1], [1, 2, 1, 0, 0], [5, 2, 5, 1, 1], [2, 1, 0, 1, 2] ], dtype=torch.float32) input = torch.reshape(input, (-1, 1, 5, 5)) print(input.shape) class Tudui(nn.Module): def __init__(self): super(Tudui, self).__init__() # ceil_mode为True,保留边框部分(多) # ceil_mode为False,不保留边框部分(少) self.maxpool = MaxPool2d(kernel_size=3, ceil_mode=False) def forward(self, input): output = self.maxpool(input) return output writer = SummaryWriter("log_maxpool") tudui = Tudui() output = tudui(input) step = 0 for data in dataLoader: imgs, targets = data writer.add_images("intput", imgs, step) output = tudui(imgs) writer.add_images("output", output, step) step = step + 1 writer.close()
线性化的实现:
import torch import torchvision from torch import nn from torch.nn import Linear from torch.utils.data import DataLoader dataset = torchvision.datasets.CIFAR10("./dataset", download=True, train=False, transform=torchvision.transforms.ToTensor()) dataLoader = DataLoader(dataset, batch_size=64) class Tudui(nn.Module): def __init__(self): super(Tudui, self).__init__() self.Linear = Linear(196608, 10) def forward(self, input): output = self.Linear(input) return output tudui = Tudui() for data in dataLoader: imgs, targets = data print(imgs.shape) output = torch.flatten(imgs) print(output.shape) output = tudui(output) print(output.shape)
loss是损失函数,pytorch中提供了几种可以直接使用的,我在这里直接举例了:
# function:损失函数使用示例 import torch from torch.nn import L1Loss from torch import nn inputs = torch.tensor([1, 2, 3], dtype=torch.float32) targets = torch.tensor([1, 2, 5], dtype=torch.float32) inputs = torch.reshape(inputs, (1, 1, 1, 3)) targets = torch.reshape(targets, (1, 1, 1, 3)) # L1直接损失 loss = L1Loss() result = loss(inputs, targets) # 均方差损失 loss_mse = nn.MSELoss() result_mes = loss_mse(inputs, targets) x = torch.tensor([0.1, 0.2, 0.3]) y = torch.tensor([1]) x = torch.reshape(x, (1, 3)) loss_cross = nn.CrossEntropyLoss() result_cross = loss_cross(x, y) print(result) print(result_mes) print(result_cross)
优化器的使用,有个模板:
optim.zero_grad() # 1.设置梯度为0
result_loss.backward() # 2.计算梯度,进行反向传播
optim.step() # 3.进行梯度更新,调整权重参数(降低loss)
示例如下:
import torch import torchvision from torch import nn from torch.nn import Conv2d, Linear, Sequential from torch.nn import MaxPool2d from torch.nn import Flatten from torch.utils.data import DataLoader, dataloader from nn_loss import loss class Tudui(nn.Module): def __init__(self): super(Tudui, self).__init__() self.model = Sequential( Conv2d(3, 32, 5, padding=2), MaxPool2d(2), Conv2d(32, 32, 5, padding=2), MaxPool2d(2), Conv2d(32, 64, 5, padding=2), MaxPool2d(2), Flatten(), Linear(1024, 64), Linear(64, 10) ) def forward(self, input): x = self.model(input) return x dataset = torchvision.datasets.CIFAR10("./dataset", train=False, transform=torchvision.transforms.ToTensor(), download=True) dataLoader = DataLoader(dataset, batch_size=1) loss = nn.CrossEntropyLoss() tudui = Tudui() optim = torch.optim.SGD(tudui.parameters(), lr=0.01) for epoch in range(20): running_loss = 0.0 for data in dataLoader: imgs, targets = data outputs = tudui(imgs) result_loss = loss(outputs, targets) optim.zero_grad() # 1.设置梯度为0 result_loss.backward() # 2.计算梯度,进行反向传播 optim.step() # 3.进行梯度更新,调整权重参数(降低loss) running_loss = running_loss + result_loss print(running_loss)
引入非线性(线性的表彰不好),示例如下:
import torch import torchvision from torch import nn from torch.nn import ReLU from torch.utils.data import DataLoader from torch.utils.tensorboard import SummaryWriter input = torch.tensor([[1, 0.5], [-1, 3]]) input = torch.reshape(input, (-1, 1, 2, 2)) print(input) dataset = torchvision.datasets.CIFAR10("./dataset", download=True, train=False, transform=torchvision.transforms.ToTensor()) dataLoader = DataLoader(dataset, batch_size=64) class Tudui(nn.Module): def __init__(self): super(Tudui, self).__init__() self.relu = ReLU() def forward(self, input): output = self.relu(input) return output writer = SummaryWriter("logs_relu") tudui = Tudui() output = tudui(input) step = 0 for data in dataLoader: imgs, targets = data writer.add_images("input", imgs, step) output = tudui(imgs) writer.add_images("output", output, step) step = step + 1 writer.close() print(output)
传统方式写模型,要在forward中一层一层写 结构。使用sequential可以直接在model中进行定义,示例如下:
## 传统方式:自己写模型 # import torch # from torch import nn # from torch.nn import Conv2d, Linear, Sequential # from torch.nn import MaxPool2d # from torch.nn import Flatten # # class Tudui(nn.Module): # def __init__(self): # super(Tudui, self).__init__() # self.conv1 = Conv2d(3, 32, 5, padding=2) # self.maxpool1 = MaxPool2d(2) # self.conv2 = Conv2d(32, 32, 5, padding=2) # self.maxpool2 = MaxPool2d(2) # self.conv3 = Conv2d(32, 64, 5, padding=2) # self.maxpool3 = MaxPool2d(2) # self.flatten = Flatten() # self.linear1 = Linear(1024, 64) # self.linear2 = Linear(64, 10) # # def forward(self, input): # x = self.conv1(input) # x = self.maxpool1(x) # x = self.conv2(x) # x = self.maxpool2(x) # x = self.conv3(x) # x = self.maxpool3(x) # x = self.flatten(x) # x = self.linear1(x) # x = self.linear2(x) # return x # # tudui = Tudui() # input = torch.ones((64, 3, 32, 32)) # output = tudui(input) # print(output.shape) ## 使用Sequential写模型 import torch from torch import nn from torch.nn import Conv2d, Linear, Sequential from torch.nn import MaxPool2d from torch.nn import Flatten from torch.utils.tensorboard import SummaryWriter class Tudui(nn.Module): def __init__(self): super(Tudui, self).__init__() self.model = Sequential( Conv2d(3, 32, 5, padding=2), MaxPool2d(2), Conv2d(32, 32, 5, padding=2), MaxPool2d(2), Conv2d(32, 64, 5, padding=2), MaxPool2d(2), Flatten(), Linear(1024, 64), Linear(64, 10) ) def forward(self, input): x = self.model(input) return x writer = SummaryWriter("logs_sequential") tudui = Tudui() input = torch.ones((64, 3, 32, 32)) output = tudui(input) writer.add_graph(tudui, input) writer.close() print(output.shape)
示例如下:
import torchvision from torch.utils.data import DataLoader import torch from torch import nn from torch.utils.tensorboard import SummaryWriter train_data = torchvision.datasets.CIFAR10("./dataset", train=True, download=True, transform=torchvision.transforms.ToTensor()) test_data = torchvision.datasets.CIFAR10("./dataset", train=False, download=True, transform=torchvision.transforms.ToTensor()) # length train_data_size = len(train_data) test_data_size = len(test_data) print("训练数据集长度为:{}".format(train_data_size)) print("测试数据集长度为:{}".format(test_data_size)) # 利用dataloader来加载数据集 train_dataloader = DataLoader(train_data, batch_size=64) test_dataloader = DataLoader(test_data, batch_size=64) # 搭建神经网络 class Tudui(nn.Module): def __init__(self): super(Tudui, self).__init__() self.model = nn.Sequential( nn.Conv2d(3, 32, 5, 1, 2), nn.MaxPool2d(2), nn.Conv2d(32, 32, 5, 1, 2), nn.MaxPool2d(2), nn.Conv2d(32, 64, 5, 1, 2), nn.MaxPool2d(2), nn.Flatten(), nn.Linear(64*4*4, 64), nn.Linear(64, 10)) def forward(self, input): x = self.model(input) return x ### 测试模型正确性 # if __name__ == '__main__': # tudui = Tudui() # input = torch.ones((64, 3, 32, 32)) # output = tudui(input) # print(output.shape) # 创建网络模型 tudui = Tudui() # 损失函数 loss_fn = nn.CrossEntropyLoss() # 优化器 learning_rate = 1e-2 optimizer = torch.optim.SGD(tudui.parameters(), lr=learning_rate) # 设置训练网络的一些参数 total_train_step = 0 total_test_step = 0 # 训练轮数 epoch = 10 # 添加tensorboard writter = SummaryWriter("logs_train") total_test_step = 0 for i in range(10): print("-----第{}轮训练开始----".format(i+1)) # 训练步骤开始 tudui.train() #特定层:Dropout for data in train_dataloader: # 1. 数据导入 imgs, targets = data # 2. 模型导入 outputs = tudui(imgs) # 3. loss计算 loss = loss_fn(outputs, targets) # 4. 优化器优化模型 optimizer.zero_grad() loss.backward() optimizer.step() total_train_step = total_train_step + 1 if total_train_step % 100 == 0: print("训练次数:{}, loss:{}".format(total_train_step, loss.item())) #loss.item()相当于取值 writter.add_scalar("train_loss", loss.item(), total_train_step) # 测试步骤开始 tudui.eval() total_test_loss = 0 total_accuracy = 0 with torch.no_grad(): for data in test_dataloader: imgs, targets = data outputs = tudui(imgs) loss = loss_fn(outputs, targets) total_test_loss = total_test_loss + loss.item() accuracy = (outputs.argmax(1) == targets).sum() total_accuracy = total_accuracy + accuracy print("整体测试集上的loss:{}".format(total_test_loss)) print("整体测试集上的正确率:{}".format(total_accuracy/test_data_size)) writter.add_scalar("test_loss", loss.item(), total_test_step) writter.add_scalar("test_accuracy", total_accuracy/test_data_size, total_test_step) total_test_step = total_test_step + 1 # 保存训练模型 torch.save(tudui, "tudui_{}.pth".format(i)) print("模型已保存!") writter.close()
示例如下:
# 与CPU的区别:在网络模型、数据、损失函数上增加cuda() import torchvision from torch.utils.data import DataLoader import torch from torch import nn from torch.utils.tensorboard import SummaryWriter train_data = torchvision.datasets.CIFAR10("./dataset", train=True, download=True, transform=torchvision.transforms.ToTensor()) test_data = torchvision.datasets.CIFAR10("./dataset", train=False, download=True, transform=torchvision.transforms.ToTensor()) # length train_data_size = len(train_data) test_data_size = len(test_data) print("训练数据集长度为:{}".format(train_data_size)) print("测试数据集长度为:{}".format(test_data_size)) # 利用dataloader来加载数据集 train_dataloader = DataLoader(train_data, batch_size=64) test_dataloader = DataLoader(test_data, batch_size=64) # 搭建神经网络 class Tudui(nn.Module): def __init__(self): super(Tudui, self).__init__() self.model = nn.Sequential( nn.Conv2d(3, 32, 5, 1, 2), nn.MaxPool2d(2), nn.Conv2d(32, 32, 5, 1, 2), nn.MaxPool2d(2), nn.Conv2d(32, 64, 5, 1, 2), nn.MaxPool2d(2), nn.Flatten(), nn.Linear(64*4*4, 64), nn.Linear(64, 10)) def forward(self, input): x = self.model(input) return x ### 测试模型正确性 # if __name__ == '__main__': # tudui = Tudui() # input = torch.ones((64, 3, 32, 32)) # output = tudui(input) # print(output.shape) # 创建网络模型 tudui = Tudui() tudui = Tudui().cuda() # 损失函数 loss_fn = nn.CrossEntropyLoss() loss_fn = loss_fn.cuda() # 优化器 learning_rate = 1e-2 optimizer = torch.optim.SGD(tudui.parameters(), lr=learning_rate) # 设置训练网络的一些参数 total_train_step = 0 total_test_step = 0 # 训练轮数 epoch = 10 # 添加tensorboard writter = SummaryWriter("logs_train") total_test_step = 0 for i in range(10): print("-----第{}轮训练开始----".format(i+1)) # 训练步骤开始 tudui.train() #特定层:Dropout for data in train_dataloader: # 1. 数据导入 imgs, targets = data imgs = imgs.cuda() targets = targets.cuda() # 2. 模型导入 outputs = tudui(imgs) # 3. loss计算 loss = loss_fn(outputs, targets) # 4. 优化器优化模型 optimizer.zero_grad() loss.backward() optimizer.step() total_train_step = total_train_step + 1 if total_train_step % 100 == 0: print("训练次数:{}, loss:{}".format(total_train_step, loss.item())) #loss.item()相当于取值 writter.add_scalar("train_loss", loss.item(), total_train_step) # 测试步骤开始 tudui.eval() total_test_loss = 0 total_accuracy = 0 with torch.no_grad(): for data in test_dataloader: imgs, targets = data imgs = imgs.cuda() targets = targets.cuda() outputs = tudui(imgs) loss = loss_fn(outputs, targets) total_test_loss = total_test_loss + loss.item() accuracy = (outputs.argmax(1) == targets).sum() total_accuracy = total_accuracy + accuracy print("整体测试集上的loss:{}".format(total_test_loss)) print("整体测试集上的正确率:{}".format(total_accuracy/test_data_size)) writter.add_scalar("test_loss", loss.item(), total_test_step) writter.add_scalar("test_accuracy", total_accuracy/test_data_size, total_test_step) total_test_step = total_test_step + 1 # 保存训练模型 torch.save(tudui, "tudui_{}.pth".format(i)) print("模型已保存!") writter.close()
示例如下:
# function:训练模型的保存
import torch
import torchvision
vgg16 = torchvision.models.vgg16(pretrained=False)
# 保存方式1
torch.save(vgg16, "vgg16_method1.pth")
# 保存方式2(官方推荐)
torch.save(vgg16.state_dict(), "vgg16_method2.pth")
示例如下:
# function:现有模型加载
import torch
import torchvision
model = torch.load("vgg16_method1.pth")
print(model)
vgg16 = torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(torch.load("vgg16_method2.pth"))
# model = torch.load("vgg16_method2.pth") # 直接是字典权重
print(model)
示例如下:
# function:自找图片,验证train.py训练的模型准确性 import torch import torchvision from PIL import Image from torch import nn image_path = "dog.png" image = Image.open(image_path) transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32, 32)), torchvision.transforms.ToTensor()]) image = transform(image) class Tudui(nn.Module): def __init__(self): super(Tudui, self).__init__() self.model = nn.Sequential( nn.Conv2d(3, 32, 5, 1, 2), nn.MaxPool2d(2), nn.Conv2d(32, 32, 5, 1, 2), nn.MaxPool2d(2), nn.Conv2d(32, 64, 5, 1, 2), nn.MaxPool2d(2), nn.Flatten(), nn.Linear(64*4*4, 64), nn.Linear(64, 10)) def forward(self, input): x = self.model(input) return x model = torch.load("tudui_9.pth") image = torch.reshape(image, (1, 3, 32, 32)) image = image.cuda() model.eval() with torch.no_grad(): output = model(image) print(output.argmax(1))
示例如下:
# function:使用现有网络对现有数据集进行训练 import torchvision from torch import nn vgg16_false = torchvision.models.vgg16(pretrained=False) vgg16_true = torchvision.models.vgg16(pretrained=True) print(vgg16_true) train_data = torchvision.datasets.CIFAR10("./dataset", train=True, download=True, transform=torchvision.transforms.ToTensor()) # CIFAR10最终的输出结果是10类,所以必须按照原来的增加一层 vgg16_true.classifier.add_module('add_linear', nn.Linear(1000, 10)) print(vgg16_true) # CIFAR10最终的输出结果是10类,也可以在原来基础上做改动 vgg16_false.classifier[6] = nn.Linear(4096, 10) print(vgg16_false)
以上简单的了解了下pytorch的基础,学习仍在继续,继续加油~
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。