本文参考
PyTorch 优点:
PyTorch VS TensorFlow
# 创建一个 pytorch 虚拟环境
conda create -n pytorch anaconda
conda activate pytorch
conda env list
# 使用国内的镜像站安装
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch/
# 安装cpu版本
conda install pytorch torchvision cpuonly
# 验证安装成功
>>> import torch, torchvision
>>> x = torch.rand(5, 3)
>>> x
tensor([[0.9567, 0.8210, 0.7273],
[0.7107, 0.1292, 0.4716],
[0.2321, 0.1172, 0.6120],
[0.0253, 0.9240, 0.5838],
[0.4741, 0.6674, 0.0979]])
CUDA
nvidia-smi
验证驱动是否装好。之后再安装 CUDA 与 cuDNN 库)conda install pytorch torchvision cudatoolkit=10.2
# 验证 GPU 驱动和 CUDA 可用
>>> import torch
>>> torch.cuda.is_available()
ndarrays
,同时 Tensor 可以使用 GPU 进行计算import torch
torch.Tensor
是默认的 tensor 类型(torch.FlaotTensor
)的简称。设置默认的 Tensor 类型 (可以在局部使用完后再重新设置回 torch.FloatTensor
):>>> torch.set_default_tensor_type('torch.DoubleTensor')
>>> a = torch.randn(2,2)
>>> a
tensor([[ 1.7017, -0.1865],
[-1.1735, 0.4421]])
int()
, float()
, double()
int()
, float()
, double()
等直接进行数据类型转换>>> b = a.double()
>>> b
tensor([[ 1.7017, -0.1865],
[-1.1735, 0.4421]], dtype=torch.float64)
type()
>>> c = a.type(torch.DoubleTensor)
>>> c
tensor([[ 1.7017, -0.1865],
[-1.1735, 0.4421]], dtype=torch.float64)
type_as()
type_as()
函数可以保持 Tensor 之间的类型一致>>> d = a.type_as(b)
>>> d
tensor([[ 1.7017, -0.1865],
[-1.1735, 0.4421]], dtype=torch.float64)
>>> torch.Tensor(2, 2)
tensor([[0.0000e+00, 0.0000e+00],
[1.6685e-10, 4.5909e-41]])
>>> torch.DoubleTensor(2, 2)
tensor([[1.6022e-306, 2.9644e-323],
[4.9407e-324, 1.1126e-306]], dtype=torch.float64)
>>> torch.FloatTensor([[1, 2, 3], [4, 5, 6]])
tensor([[1., 2., 3.],
[4., 5., 6.]])
>>> torch.Tensor([[1, 2, 3], [4, 5, 6]])
tensor([[1., 2., 3.],
[4., 5., 6.]])
>>> torch.Tensor(np.array([[1, 2, 3], [4, 5, 6]]))
tensor([[ 1, 2, 3],
[ 4, 5, 6]])
>>> torch.tensor([[1., -1.], [1., -1.]])
tensor([[ 1.0000, -1.0000],
[ 1.0000, -1.0000]])
>>> torch.tensor(np.array([[1, 2, 3], [4, 5, 6]]))
tensor([[ 1, 2, 3],
[ 4, 5, 6]])
torch.Tensor()
always copies data. If you have a Tensor data and just want to change itsrequires_grad
flag, userequires_grad_()
ordetach()
to avoid a copy. If you have a numpy array and want to avoid a copy, usetorch.as_tensor()
.
torch.zeros
, torch.ones
, torch.eye
>>> torch.zeros(5, 3, dtype=torch.long)
tensor([[0, 0, 0],
[0, 0, 0],
[0, 0, 0],
[0, 0, 0],
[0, 0, 0]])
>>> torch.ones(2, 2)
tensor([[1., 1.],
[1., 1.]])
# 生成 对角矩阵
>>> torch.eye(2, 2)
tensor([[1., 0.],
[0., 1.]])
>>> torch.randn(2, 2)
tensor([[-1.3644, 0.9027],
[ 0.3454, -0.6466]])
# randperm(num)生成长度为num的随机排列向量
>>> torch.randperm(10)
tensor([1, 6, 0, 5, 2, 8, 7, 3, 9, 4])
torch.arange
,torch.linspace
>>> torch.arange(1, 6, 2)
tensor([1, 3, 5])
>>> torch.linspace(1, 6, 2)
tensor([1., 6.])
torch.randn_like
,torch.ones_like
,torch.zeros_like
>>> x
tensor([[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.]], dtype=torch.float64)
>>> x = torch.randn_like(x, dtype=torch.float)
>>> x
tensor([[ 0.8157, 0.6115, 0.5099],
[-1.7238, 0.7658, 0.1270],
[ 0.2794, -0.2724, 0.4168],
[ 1.5144, -1.6933, -0.6357],
[-0.1833, 1.0657, -1.0282]])
>>> torch.ones_like(x)
tensor([[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.]])
>>> torch.zeros_like(x)
tensor([[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]])
>>> x.size()
torch.Size([5, 3])
>>> x.shape
torch.Size([5, 3])
>>> x.numel()
15
>>> x.nelement()
15
torch.cat()
, torch.stack()
>>> a = torch.arange(1, 5).reshape(2, 2)
>>> b = torch.arange(5, 9).reshape(2, 2)
>>> a
tensor([[1, 2],
[3, 4]])
>>> b
tensor([[5, 6],
[7, 8]])
>>> torch.cat([a, b], 0)
tensor([[1, 2],
[3, 4],
[5, 6],
[7, 8]])
>>> torch.cat([a, b], 1)
tensor([[1, 2, 5, 6],
[3, 4, 7, 8]])
# 以第0维进行stack,叠加的基本单位为序列本身,即a和b,因此输出[a, b]
>>> torch.stack([a, b], 0)
tensor([[[1, 2],
[3, 4]],
[[5, 6],
[7, 8]]])
>>> torch.stack([a, b], 0).shape
torch.Size([2, 2, 2])
# 以第1维进行stack,叠加的基本单位为一行
>>> torch.stack([a, b], 1)
tensor([[[1, 2],
[5, 6]],
[[3, 4],
[7, 8]]])
>>> torch.stack([a, b], 1).shape
torch.Size([2, 2, 2])
# 以第2维进行stack,叠加的基本单位为每一行的每一个元素
>>> torch.stack([a, b], 2)
tensor([[[1, 5],
[2, 6]],
[[3, 7],
[4, 8]]])
>>> torch.stack([a, b], 2).shape
torch.Size([2, 2, 2])
torch.chunk()
, torch.split()
torch.chunk()
需要指定分块的数量>>> a = torch.arange(6).reshape(2, 3)
>>> a
tensor([[0, 1, 2],
[3, 4, 5]])
# 沿着第0维进行分块,一共分成两块
>>> torch.chunk(a, 2, 0)
(tensor([[0, 1, 2]]), tensor([[3, 4, 5]]))
# 沿着第1维进行分块,一共分成两块.当不能整除时,最后一个的维数会小于前面的
>>> torch.chunk(a, 2, 1)
(tensor([[0, 1],
[3, 4]]), tensor([[2],
[5]]))
torch.split()
需要指定每一块的大小,以整型或者 list
表示# 沿着第0维分块,每一块维度为2,由于第0维维度总共为2,因此相当于没有分割
>>> torch.split(a, 2, 0)
(tensor([[0, 1, 2],
[3, 4, 5]]),)
# 沿着第1维分块,每一块维度为2
>>> torch.split(a, 2, 1)
(tensor([[0, 1],
[3, 4]]), tensor([[2],
[5]]))
# 沿着第1维分块,list中的元素代表每一块占的维度
>>> torch.split(a, [1, 2], 1)
(tensor([[0],
[3]]), tensor([[1, 2],
[4, 5]]))
ndarray
,也支持布尔索引和神奇索引>>> a=torch.randn(2,2)
>>> a
tensor([[ 0.3150, -1.8297],
[ 0.0648, -1.1429]])
>>> a>0
tensor([[ True, False],
[ True, False]])
>>> a[~(a>0)]
tensor([-1.8297, -1.1429])
>>> x
tensor([[0, 0, 0, 0],
[1, 1, 1, 1],
[2, 2, 2, 2],
[3, 3, 3, 3],
[4, 4, 4, 4],
[5, 5, 5, 5],
[6, 6, 6, 6],
[7, 7, 7, 7]])
>>> x[[4, 3, 0, 6]]
tensor([[4, 4, 4, 4],
[3, 3, 3, 3],
[0, 0, 0, 0],
[6, 6, 6, 6]])
>>> x[[0, 3, 1, 2], :]
tensor([[0, 0, 0, 0],
[3, 3, 3, 3],
[1, 1, 1, 1],
[2, 2, 2, 2]])
>>> x[[0, 1, 2, 3], [0, 1, 2, 3]]
tensor([0, 1, 2, 3])
view()
, reshape()
>>> a
tensor([0, 1, 2, 3])
>>> a.view(2,2)
tensor([[0, 1],
[2, 3]])
>>> a.reshape(2,2)
tensor([[0, 1],
[2, 3]])
contiguous()
transpose()
, permute()
等可能会把 Tensor 在内存中变得不连续,而有些操作如 view()
是需要 Tensor 内存连续的,这种情况下需要使用 contiguous()
将内存变为连续的reshape()
就相当于 Tensor.contiguous().view()
resize_()
>>> a
tensor([[0, 1],
[2, 3]])
>>> a.resize_(8)
tensor([ 0, 1, 2, 3,
1910660200560, 1910660200560, 1910660200560, 1910660200560])
>>> a
tensor([ 0, 1, 2, 3,
1910660200560, 1910660200560, 1910660200560, 1910660200560])
# 如果resize后小于原Tensor大小,则剩余的部分仍然会隐藏保留
>>> a.resize_(2)
tensor([0, 1])
>>> a
tensor([0, 1])
>>> a.resize_(4)
tensor([0, 1, 2, 3])
squeeze()
, unsqueeze()
squeeze()
去除 size 为 1 的维度unsqueeze()
将指定维度的 size 变为 1>>> a=torch.arange(3)
>>> a.shape
torch.Size([3])
>>> a.unsqueeze(0).shape
torch.Size([1, 3])
>>> a.unsqueeze(0).squeeze(0).shape
torch.Size([3])
expand()
, expand_as()
expand()
将 size 为 1 的维度复制扩展为指定大小,也可以用 expand_as()
指定 Tensor 维度>>> a=torch.randn(2,2,1)
>>> a
tensor([[[ 0.1891],
[ 2.6933]],
[[-1.6266],
[ 1.3337]]])
>>> a.expand(2,2,3)
tensor([[[ 0.1891, 0.1891, 0.1891],
[ 2.6933, 2.6933, 2.6933]],
[[-1.6266, -1.6266, -1.6266],
[ 1.3337, 1.3337, 1.3337]]])
>>> b=torch.randn(2,2,3)
>>> a.expand_as(b)
tensor([[[ 0.1891, 0.1891, 0.1891],
[ 2.6933, 2.6933, 2.6933]],
[[-1.6266, -1.6266, -1.6266],
[ 1.3337, 1.3337, 1.3337]]])
transpose()
和 permute()
transpose()
将指定的两个维度的元素进行转置permute()
按照给定的维度进行维度变换, 相当于 numpy 中的 transpose()
>>> a=torch.randn(2,2,2)
>>> a
tensor([[[-0.0487, 0.1411],
[-0.9727, -1.0517]],
[[-0.6619, 0.1907],
[ 0.1305, -1.2069]]])
>>> a.transpose(0,1)
tensor([[[-0.0487, 0.1411],
[-0.6619, 0.1907]],
[[-0.9727, -1.0517],
[ 0.1305, -1.2069]]])
>>> a.permute(2,1,0)
tensor([[[-0.0487, -0.6619],
[-0.9727, 0.1305]],
[[ 0.1411, 0.1907],
[-1.0517, -1.2069]]])
sort()
sort()
沿指定维度排序,返回排序后的 Tensor 及对应的索引位置>>> a=torch.randn(3,3)
>>> a
tensor([[ 1.9710, -1.0310, -0.6732],
[-0.4766, 0.3254, 1.5063],
[-0.5229, -0.0387, 0.8007]])
# True表示降序
>>> a.sort(0, True)[0]
tensor([[ 1.9710, 0.3254, 1.5063],
[-0.4766, -0.0387, 0.8007],
[-0.5229, -1.0310, -0.6732]])
>>> a.sort(0, True)[1]
tensor([[0, 1, 1],
[1, 2, 2],
[2, 0, 0]])
max()
, min()
>>> a
tensor([[ 1.9710, -1.0310, -0.6732],
[-0.4766, 0.3254, 1.5063],
[-0.5229, -0.0387, 0.8007]])
>>> a.max(0)[0]
tensor([1.9710, 0.3254, 1.5063])
>>> a.max(0)[1]
tensor([0, 1, 1])
>>> a.max()
tensor(1.9710)
torch.max(input, other, out=None) → Tensor
input
is compared with the corresponding element of the tensor other
and an element-wise maximum is taken. The shapes of input
and other
don’t need to match, but they must be broadcastable._
”,如 add_()
, resize_()
>>> a
tensor([[1., 1.],
[1., 1.]])
>>> b
tensor([[ 0.5857, -0.7018],
[-0.2838, -0.4451]])
>>> c=a.add_(b)
# add_()使a也发生了改变
>>> a
tensor([[1.5857, 0.2982],
[0.7162, 0.5549]])
torch.as_tensor()
/ torch.from_numpy()
ndarray
来曲线救国; 使用 torch.as_tensor()
进行转换是会共享内存的>>> b=np.random.randn(2,2)
>>> a=torch.as_tensor(b) # a=torch.from_numpy(b)
>>> a
tensor([[ 0.1027, 0.0642],
[-0.1120, 0.2536]], dtype=torch.float64)
>>> b
array([[ 0.10268475, 0.0642312 ],
[-0.11195517, 0.25361295]])
>>> b[0,0]=1
>>> a
tensor([[ 1.0000, 0.0642],
[-0.1120, 0.2536]], dtype=torch.float64)
ndarray
可以使用 Tensor.numpy()
方法,该方法返回一个 numpy 数组函数名 | 功能 |
---|---|
torch.where(condition, x, y) |
满足 condition 的位置输出 x ,否则输出 y |
torch.where()
>>> x
tensor([[0, 0, 0, 0],
[1, 1, 1, 1],
[2, 2, 2, 2],
[3, 3, 3, 3],
[4, 4, 4, 4],
[5, 5, 5, 5],
[6, 6, 6, 6],
[7, 7, 7, 7]])
>>> torch.where(x > 4, torch.ones_like(x), x)
tensor([[0, 0, 0, 0],
[1, 1, 1, 1],
[2, 2, 2, 2],
[3, 3, 3, 3],
[4, 4, 4, 4],
[1, 1, 1, 1],
[1, 1, 1, 1],
[1, 1, 1, 1]])
Tensor.item()
torch.Tensor.item()
to get a Python number from a tensor containing a single value:>>> x = torch.tensor(2.5)
>>> x
tensor(2.5000)
>>> x.item()
2.5
Tensor.clamp()
Tensor.clamp(min, max)
: 限制 Tensor 元素的最大值与最小值>>> a
tensor([[0, 1],
[2, 3]])
>>> a.clamp(1, 2)
tensor([[1, 1],
[2, 2]])
Tensor.abs
Tensor.sqrt
Tensor.log
Tensor.pow
Tensor.sin
, Tensor.cos
Tensor.matmul
.requires_grad
, .grad
, .grad_fn
torch.Tensor
是包的核心类。如果将其属性 .requires_grad
(默认为 False
) 设置为 True
,则会开始跟踪针对 tensor 的所有操作,且依赖于该 Tensor 之后的所有节点都需要求导。完成计算后,可以调用 Tensor.backward()
来自动计算所有梯度。该张量的梯度将累积到 .grad
属性中Function
。Tensor 和 Function 互相连接并构建一个非循环图,它保存整个完整的计算过程的历史信息。每个张量都有一个 .grad_fn
属性保存着创建了张量的 Function 的引用,即该 Tensor 经过了怎么样的操作,用作反向传播的梯度计算(如果用户自己创建张量,则 grad_fn
是 None
)>>> x = torch.randn(2, 2, requires_grad=True)
>>> x
tensor([[-1.1423, -0.0149],
[-0.4527, -0.4478]], requires_grad=True)
>>> y = x + 2
>>> y
tensor([[0.8577, 1.9851],
[1.5473, 1.5522]], grad_fn=<AddBackward0>)
>>> y.requires_grad # 由于依赖的Tensor需要求导,因此y也需要求导
True
>>> z = 3 * y**2
>>> out = z.mean()
>>> z
tensor([[ 2.2071, 11.8215],
[ 7.1822, 7.2281]], grad_fn=<MulBackward0>)
>>> out
tensor(7.1097, grad_fn=<MeanBackward0>)
.detach_()
, .requires_grad_()
, with torch.no_grad()
.detach()
,它将其与计算历史记录分离,并防止将来的计算被跟踪。还可以将代码块使用 with torch.no_grad()
: 包装起来。在评估模型时特别有用。.requires_grad()
>>> a = torch.randn(2, 2)
>>> a.requires_grad
False
>>> a.requires_grad_()
tensor([[-0.2342, -1.2398],
[-0.8127, -0.9557]], requires_grad=True)
>>> a.detach_()
tensor([[-0.2342, -1.2398],
[-0.8127, -0.9557]])
>>> a.requires_grad
False
>>> b
tensor([1.], requires_grad=True)
>>> with torch.no_grad():
... print((b**2).requires_grad)
...
False
Tensor.backward()
Tensor.backward()
。.backward()
还有一个参数 grad_variabels
,代表根结点的导数,如果 Tensor 是标量,则不需要指定任何参数,backward()
等同于 backward(torch.tensor(1.))
。但是如果它有更多元素,则需要指定一个 gradient
参数来指定张量的形状retain_graph
设置为 True
(如果设置为 False
计算图中的中间变量在计算完后就会被释放),使计算多个输出的梯度时互不影响>>> x = torch.randn(1)
>>> x
tensor([0.2249])
>>> w = torch.ones(1, requires_grad=True)
>>> b = torch.ones(1, requires_grad=True)
# is_leaf表示x, w, b为计算图中的叶结点
>>> x.is_leaf, w.is_leaf, b.is_leaf
(True, True, True)
>>> y = w * x
>>> z = y + b
>>> z.backward(retain_graph=True)
>>> w.grad
tensor([0.2249])
>>> b.grad
tensor([1.])
torch.nn
import os
import json
from PIL import Image
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import optim
import torchvision
from torchvision import models
from torch.utils.data import Dataset
from torchvision import transforms
from torch.utils.data import DataLoader
torch.nn
接口构建于 autograd 之上,提供了网络模组、优化器和初始化策略等一系列功能nn.Module
类nn.Module
类实现了网络各层的定义及前向传播与反向传播机制
nn.Module
类,在初始化中定义模型结构和参数,在 forward()
中编写网络前向传播过程即可。nn.Module
类的 __call__
中调用了 forward
方法nn.Parameter()
Tensor
构造方式,可在 nn.Module
子类的 __init__()
中定义网络参数。默认 requires_grad=True
nn.functional
库nn.functional
库也提供了很多网络层和函数,但 nn.functional
定义的网络层不可自动学习参数,还需要使用 nn.Parameter
封装。nn.functional
的设计初衷是定义一些不需要学习参数的层,如激活层、池化层、BN 层nn.Module
,无参数学习的层,可以使用 nn.functional
class Linear(nn.Module):
def __init__(self, in_dim, out_dim):
super(Linear, self).__init__()
self.w = nn.Parameter(torch.randn(in_dim, out_dim))
self.b = nn.Parameter(torch.randn(out_dim))
def forward(self, x):
x = x.matmul(self.w)
y = x + self.b
return y
class Perception(nn.Module):
def __init__(self, in_dim, hid_dim, out_dim):
super(Perception, self).__init__()
self.layer1 = Linear(in_dim, hid_dim)
self.layer2 = Linear(hid_dim, out_dim)
def forward(self, x):
x = self.layer1(x) # nn.Module类的__call__中调用了forward方法
y = torch.sigmoid(x)
y = self.layer2(y)
y = torch.sigmoid(y)
return y
nn.Module
实例的网络结构以及参数__str__()
net.named_parameters()
, net.parameters()
net.named_parameters()
/ net.parameters()
查看网络参数nn.Sequential()
nn.Sequential()
快速搭建模型,而不必在 forward()
中一层层地前向传播nn.Sequential()
实现多层感知机class MLP(nn.Module):
def __init__(self, in_dim, hid_dim1, hid_dim2, out_dim):
super(MLP, self).__init__()
self.layer = nn.Sequential(
nn.Linear(in_dim, hid_dim1),
nn.ReLU(),
nn.Linear(hid_dim1, hid_dim2),
nn.ReLU(),
nn.Linear(hid_dim2, out_dim),
nn.Softmax(dim=1)
)
def forward(self, x):
# 直接调用__call__
x = self.layer(x)
# 也可以利用迭代协议
for layer in self.layer:
x = layer(x)
return x
nn.Sequential()
动态添加网络层torch.nn.Module.add_module(name, module)
name
(string
) – name of the child module. The child module can be accessed from this module using the given namemodule
(Module
) – child module to be added to the module.net = nn.Sequential(nn.Linear(100, 10))
net.add_module('extra-layer',nn.Linear(10, 20))
Sequential(
(0): Linear(in_features=100, out_features=10, bias=True)
(1): Linear(in_features=10, out_features=20, bias=True)
)
torch.nn
和 torch.nn.functional
中都提供了各种损失函数,它们的功能是一样的,区别在于 torch.nn
将损失函数看作网络的某一层放到模型定义中,而 torch.nn.functional
则是直接将损失函数当作功能函数放到前向传播的过程中# 以一个样本的三分类问题为例
>>> y = torch.Tensor([0.1, 0.3, 0.6]).unsqueeze(0)
>>> t = torch.Tensor([2]).long()
>>> loss_f = F.cross_entropy(y, t)
>>> loss_f
tensor(0.8533)
>>> criterion = nn.CrossEntropyLoss()
>>> loss_nn = criterion(y, t)
>>> loss_nn
tensor(0.8533)
nn.optim
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
# 每次优化之前都要先清空梯度
optimizer.zero_grad()
loss.backward()
optimizer.step()
学习率调整:
# 手动调整学习率
def set_learning_rate(optimizer, lr):
"""Sets the learning rate to the given value"""
for param_group in optimizer.param_groups:
param_group['lr'] = lr
torchvision.models
库torchvision.models
库提供了众多经典网络结构与预训练模型,如 VGG, ResNet, Inceptionvgg.features
与 vgg.classifier
表示,每个部分都是 nn.Sequential
实例>>> from torchvision import models
>>> net = models.vgg16()
# VGG16的特征层包含13个卷积层、13个Relu层、5个池化层
>>> len(net.features)
31
# VGG16的分类层包含3个全连接层、2个Relu层、2个Dropout层
>>> len(net.classifier)
7
# 可以通过索引选取网络中的某些层
>>> net.classifier[-1]
Linear(in_features=4096, out_features=1000, bias=True)
torchvision.models
中的预训练模型net = models.vgg16(pretrained=True)
model = models.vgg16(pretrained=False)
state_dict = torch.load('./vgg16.pth')
model.load_state_dict(state_dict )
torch.save()
state_dict()
方法可以获取网络的当前状态数据# 以VGG16为例,net.state_dict()中保存了该网络16个带权重的层的参数
>>> type(net.state_dict())
<class 'collections.OrderedDict'>
>>> for key in net.state_dict():
... print(key)
...
features.0.weight
features.0.bias
features.2.weight
features.2.bias
features.5.weight
features.5.bias
features.7.weight
features.7.bias
features.10.weight
features.10.bias
features.12.weight
features.12.bias
features.14.weight
features.14.bias
features.17.weight
features.17.bias
features.19.weight
features.19.bias
features.21.weight
features.21.bias
features.24.weight
features.24.bias
features.26.weight
features.26.bias
features.28.weight
features.28.bias
classifier.0.weight
classifier.0.bias
classifier.3.weight
classifier.3.bias
classifier.6.weight
classifier.6.bias
torch.save({
'model': net.state_dict(),
'optimizer': optimizer.state_dict()},
'D:/vgg16.pth')
model.load_state_dict()
net = models.vgg16()
state_dict = torch.load('D:/vgg16.pth')
net.load_state_dict({k:v for k, v in state_dict['model'].items() if k in net.state_dict()})
for layer in range(7):
for p in net.features[layer].parameters():
p.requires_grad = False
torchvision.datasets
直接加载 Imagenet,CIFAR10,MNIST 等公共数据集; 也可以继承 torch.utils.data.Dataset
抽象类,实现 __len__()
和 __getitem__()
方法,即可进行数据集迭代
__len__()
用来提供数据集大小(可选),而 __getitem__()
用来支持整数索引 (必须)torchvision.transforms
torchvision.transforms
可以方便的进行图像缩放、裁剪、随机翻转、填充以及张量的归一化等操作。操作对象可以是 PIL 的 Image 或者 Tensortransforms.Compose
将多个变换整合起来。在实际使用时,常会将变换操作集成到 Dataset
类中torch.utils.data.DataLoader
类
torch.utils.data.DataLoader
类包装之后就可以实现批量处理、随机选取等操作DataLoader
类是一个可迭代对象,对它的实例进行迭代即可用于训练过程img
和 json
文件夹中分别存有图片与标注,文件名为 0.jpg
, 1.jpg
… 0.json
, 1.json
… 读取代码如下:class MyData(Dataset):
def __init__(self, img_path, annotation_path, transforms=None):
# 初始化,读取数据集
self.annotation_path = annotation_path
self.img_path = img_path
self.transforms = transforms
def __len__(self):
return len(os.listdir(self.img_path))
def __getitem__(self, index):
annotation = json.load(open(self.annotation_path + '/' + str(index) + '.json'))
img = Image.open(self.img_path + '/' + str(index) + '.jpg')
# plt.imshow(img)
if self.transforms:
img = self.transforms(img)
return img, annotation
dataset = MyData('D:/Download/Dataset/traffic_light/train/img', 'D:/Download/Dataset/traffic_light/train/json',
transforms=transforms.Compose([
transforms.Resize(240), # 将图像最短边缩至240,宽高比例不变
transforms.RandomHorizontalFlip(), # 以0.5的概率左右翻转图像
transforms.ToTensor(), # 将PIL图像转为Tensor,并且进行归一化
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) # 进行mean与std为0.5的标准化
]))
dataloader = DataLoader(dataset, batch_size=4, shuffle=True)#, num_workers=2) # num_workers表示使用几个线程来加载数据 我的电脑加了这个参数就报错,可能不支持多线程操作
data_iter = iter(dataloader)
for step in range(1000):
data = next(data_iter)
# 下面即可将data用于训练网络
torch.cuda.is_available
判断当前 GPU 是否可用cuda()
方法可以将数据转移到 GPU 上运行,并且可以输入数字来指定转移到哪一块 GPU 上运行a = torch.randn(3, 3)
b = models.vgg16()
if torch.cuda.is_available():
a = a.cuda()
b = b.cuda(1)
torch.device()
指定使用哪一个GPUdevice = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
c = torch.randn(3, 3, device=device)
model.to(device)
c = c.to(device)
# 在终端执行脚本时指定
CUDA_VISIBLE_DEVICES=2 python train.py
torch.cuda.set_device(1)
model_gpu = nn.DataParallel(model)
output = model_gpu(input)