本文为 PyTorch 学习笔记，介绍了 PyTorch 项目的基本代码范式。
数据的导入
数据的导入需要用到两个 PyTorch 库，分别是 torch.utils.data.Dataset 和 torch.utils.data.DataLoader
Dataset 类的使用
可以使用 Dataset 类处理自定义的数据集，示例如下：
class MyDataset(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.image_name = os.listdir(self.path)
def __getitem__(self, idx):
img_name = self.image_name[idx]
img_path = os.path.join(self.path, img_name)
label = self.label_dir
img = Image.open(img_path)
return img, label
def __len__(self):
return len(self.image_name)
自定义自己的数据类需要继承 PyTorch 的 Dataset 类。
方法的重写
继承 Dataset 父类后，还需要根据数据集的具体结构和任务需求重写一些方法，其中必须重写的方法包括：__init__ 、__getitem__、__len__
构造方法的重写
重写构造方法的目的是定义其他方法所需的属性，例如上文示例中的根路径、标签路径等等。
getitem 方法的重写
除了对象 self 以外，__getitem__ 方法接受一个额外的变量 idx，表示数据集中一个数据的索引。
我们需要重写该方法，根据传入的索引返回将要输入模型的特征和标签。
len 方法的重写
重写该方法以定义数据集的长度，例如上文示例用图片数量定义数据集长度
使用已有数据集
这些提供的数据集使用方式很简单，只需要按照 PyTorch 文档中的使用方法填写参数，只需一行代码即可。
例如 CIFAR10 的使用方法如下：
test_set = torchvision.datasets.CIFAR10(
"./dataset",
train=False,
transform=torchvision.transforms.ToTensor(),
download=True
)
数据的处理
处理图片
PyTorch 在 torchvision.transforms 中提供了许多处理图片的工具，例如常用的工具类 torchvision.transforms.ToTensor() 用于将 PIL 或 ndarray (使用 opencv) 格式的图片转换为 PyTorch 的 Tensor 类型，以便输入模型。
除此之外，transform 模块还提供了常见的图片处理工具，例如图片反转、切割等等。
DataLoader 类的使用
test_loader = DataLoader(dataset=test_set, batch_size=64, shuffle=True,
num_workers=0, drop_last=False)
遍历数据集的方法：
for data in test_loader:
imgs, targets = data
搭建模型
搭建模型有两种方式，分别是：
- 从零搭建
- 从已有的模型搭建
从零搭建模型
首先，定义一个继承 nn.module 的模型类，然后重写方法即可，示例如下：
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.model = nn.Sequential(
nn.Conv2d(3, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.LeakyReLU(),
nn.Conv2d(32, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.LeakyReLU(),
nn.Conv2d(32, 64, 5, 1, 2),
nn.MaxPool2d(2),
nn.LeakyReLU(),
nn.Flatten(),
nn.Linear(64 * 4 * 4, 256),
nn.ReLU(),
nn.Linear(256, 64),
nn.ReLU(),
nn.Linear(64, 10),
)
def forward(self, x):
x = self.model(x)
return x
从已有的模型搭建模型
导入已有模型
获取已有的模型有两种方式，第一种方式是使用 PyTorch 提供的知名模型，例如 vgg16。
vgg16_pt = torchvision.models.vgg16(pretrained=True)
参数 pretrained 为 True 时会下载预训练后的模型参数，反之则使用未训练的随机模型参数
修改已有模型
有时候，下载的模型无法满足我们的任务需求，例如处理分类问题时，模型的输出层的神经元个数与我们需要区分的类别数量不同。这个时候，我们就需要修改模型。
修改模型所用到的函数或方法主要有以下几个：
# 1. 增加：add_module 方法，可以增加单层也可以增加 Sequential
vgg16.add_module("mod", nn.Sequential(OrderedDict([ # 可以不给每层取名
("linear", nn.Linear(1000, 256)),
("relu", nn.ReLU()),
("softmax", nn.Softmax(10))
])))
vgg16.add_module(nn.Linear(64, 10))
# 2. 修改：直接使用 [idx] 和 `.`
print(vgg16) # 先print查看模型结构
vgg16.classifier[6] = nn.Linear(64, 10)
vgg16.classifier[6] = nn.Sequential(
nn.Linear(64, 10),
nn.Linear(10, 5)
)
# 3. 删除：使用 del 即可
del vgg16.classifier
训练模型
下面提供一个训练示例：
# * 3. Create the model object and Setting hyperparameters
# create model object
model = Model()
if torch.cuda.is_available():
model = model.cuda()
# define the loss function
loss_fn = nn.CrossEntropyLoss()
if torch.cuda.is_available():
loss_fn = loss_fn.cuda()
# define the optimizer
learning_rate = 0.01
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
# optimizer = torch.optim.Adam(model.parameters())
# Setting
total_train_step = 0
total_test_step = 0
epoch = 30
# set tensorboard
writer = SummaryWriter("./logs_model")
# * 4. Training the model
for i in range(epoch):
print(f"------ Epoch {i} begin. ------")
start_time = time.time()
# train
model.train() # for some certain layer(such as dropout), See PyTorch.org.
for data in train_dataloader:
imgs, targets = data
# move to gpu
if torch.cuda.is_available():
imgs, targets = imgs.cuda(), targets.cuda()
output = model(imgs)
loss = loss_fn(output, targets)
# optimize
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_train_step = total_train_step + 1
if total_train_step % 100 == 0:
print(f"Train step: {total_train_step} / Loss: {loss}")
end_time = time.time()
print(f"tic {end_time-start_time}")
writer.add_scalar("train_loss(step)", loss, total_train_step)
# test
model.eval() # for some certain layer
total_test_loss = 0
total_accuracy = 0
with torch.no_grad():
for data in test_dataloader:
imgs, targets = data
# move to gpu
if torch.cuda.is_available():
imgs, targets = imgs.cuda(), targets.cuda()
output = model(imgs)
# compute accuracy
accuracy = (output.argmax(1) == targets).sum()
total_accuracy += accuracy
loss = loss_fn(output, targets)
total_test_loss += loss
# Show and log test info
print(ColorText.info(f"\nTotal loss in test set: {total_test_loss}"))
print(
ColorText.info(f"Total accuracy in test set: {total_accuracy / test_set_size}")
)
writer.add_scalar("test_loss(epoch)", total_test_loss, total_test_step)
writer.add_scalar(
"test_accuracy(epoch)", total_accuracy / test_set_size, total_test_step
)
total_test_step += 1
# Save Model
torch.save(model, f"model_{i}.pt")
print(ColorText.info(f"..::Model {i} Saved.."))
writer.close()
附录
Tensorboard 的使用
示例：
首先在脚本内记录数据到指定路径
from torch.urils.tensorboard import SummaryWriter
# writer 对象
writer = SummaryWriter("./logs") # log_dir
# 记录数据
writer.add_scalar("loss curve", y_value, x_value)
#
writer.close()
然后，命令行用 Tensorboard 打开日志路径
tensorboard --log_dir ./logs
评论