数据处理

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
root_dir = "xxx"
train_dir = "xxx"

class Dataset(object):
    def __init__(self, path, transform = None):
        self.xxx = xxx
        # 获取所有图片路径
        self.image_paths = xxx
        # 获取所有图片的标签
        self.labels = xxx
        ...
    def __len__(self):
        return len(self.inputs)
    
    def __getitem__(self, idx)
        # 对图像进行处理
        image = Image.read(self.image_paths)  # 读取图像
        if transform is not None:
            image = transform(image)  # 图像变换
        # 对标签进行处理
        label = self.labels[idx]
        ...
        # 返回元组  
        return image,label

模型

1
2
3
4
5
6
7
8
9
10
11
12
class Net(nn.Module):
    def __init__(self):        
        super(ConvRnn,self).__init__()
        # 定义层
        self.layer_1 = nn.Sequential(...)
        self.layer_2 = nn.Sequential(...)
        ...
    def forward(self, X, ...):
        y = self.layer_1(X)
        y = self.layer_2(y)
        ...
        return y

训练

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

# 运行在 cpu 还是 gpu
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

# 定义超参数
epoch = 32
batch_size = 64
lr = 0.01
...

# 定义迭代器
trans = transforms.Compose([transforms.ToTensor(),
			 transforms.Resize(256),
			 transform.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))],
             ...) 

dataset=Dataset(...)
dataloader=DataLoader(dataset,batch_size=batch_size,shuffle=True,num_workers=0)

# 定义模型、损失函数、优化器...
net = Net().to(device)
loss = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(net.parameters(), lr=lr)

for epoch in range(epoch):
    for images,labels in dataloader:
        y_hat = net(images)
        # 计算损失
        l = loss(y_hat,labels)
        # 反向传播
        l.backward()
        # 优化
        optimizer.step()
        ...
    # 保存模型
    torch.save(net.state_dict(), f"./checkpoints/{epoch}_weights.pth")

测试

1
2
3
4
5
6
7
8
9
10

net = Net(...)
model.load_state_dict(torch.load(...))
model.eval()

image = Image.read(...)
image = trans(image)  # 转为 tensor
image.unsqueeze_(0)  # 增加 batch 的维度

y_hat = model(image)