自上一题文本分类,这道基于预训练模型的图像分类更贴合我的科研任务。因此这道题我采用Pytorch框架进行代码编写。
题目要求如下:
这道题我就采用Torch进行编写
import torch
import json
from torch.utils.data import DataLoader, Dataset
import torch.nn as nn
import numpy as np
import csv
# Solve the dataset,取出三个json文件中的feature
with open('/home/project/resnet_train.json', 'r', encoding='utf-8') as file:
resnet_json = json.load(file)
id_list = resnet_json.keys()
with open('/home/project/inception_train.json', 'r', encoding='utf-8') as file:
inception_json = json.load(file)
with open('/home/project/xception_train.json', 'r', encoding='utf-8') as file:
xception_json = json.load(file)
class Dataset_loader(Dataset): # 自己构建一个加载器,完成feature和label的封装,注意这个类一定要继承Dataset
def __init__(self, id, resnet_data, inception_data, xception_data):
super().__init__()
self.id = id
self.resnet = resnet_data
self.inception = inception_data
self.xception = xception_data
self.feature, self.labels = [], []
for key in self.id:
temp = self.resnet[key]['feature'] + self.inception[key]['feature'] + self.xception[key]['feature']
temp = torch.tensor(temp)
self.feature.append(temp)
self.labels.append(torch.tensor(self.resnet[key]['label'], dtype=torch.float32))
def __len__(self):
return len(self.id)
def __getitem__(self, index):
return self.feature[index], self.labels[index]
train_data = Dataset_loader(id_list, resnet_json, inception_json, xception_json)
train_loader = DataLoader(train_data, batch_size=32, shuffle=True) # 转换为迭代器进行训练
class Classifier_Net(nn.Module): # 构建的网络为全Linear的结构, 大致为 2048 * 3 -》 2048 -》 1,本题还是二分类
def __init__(self, input_size, output_size, hidden_size=None):
super().__init__()
self.hidden_size = [input_size, input_size // 3, input_size // 6,
output_size] if hidden_size is None else hidden_size
self.layer = nn.Sequential(
*[nn.Sequential(nn.Linear(self.hidden_size[idx], self.hidden_size[idx + 1], bias=True), nn.GELU()) for idx
in range(len(self.hidden_size) - 2)]
)
self.predict_layer = nn.Linear(self.hidden_size[-2], self.hidden_size[-1])
def forward(self, feature):
x = self.layer(feature)
return self.predict_layer(x)
labels = []
for key in id_list:
labels.append(resnet_json[key]['label'])
class_num = np.unique(labels).shape[0]
model = Classifier_Net(input_size=2048 * 3, output_size=1)
train_parameter = {
'learning_rate': 2e-4,
'epochs': 10
}
optim_loader = torch.optim.Adam(params=model.parameters(), lr=train_parameter['learning_rate'])
criterion = torch.nn.BCEWithLogitsLoss(reduction='mean')
for epoch in range(train_parameter['epochs']):
epoch_loss = 0
for feature, label in train_loader:
optim_loader.zero_grad()
output = model(feature).squeeze()
batch_loss = criterion(output, label)
batch_loss.backward()
optim_loader.step()
epoch_loss += batch_loss.item()
print(f"Now Ending Training of epoch{epoch}, epoch loss is {epoch_loss}")
model.eval()
with open('/home/project/test.json', 'r', encoding='utf-8') as file:
test_json = json.load(file)
test_id_list = test_json.keys() # 加载索引以及特征
with open('/home/project/result.csv', 'w', newline='') as file:
csv_writer = csv.writer(file)
csv_writer.writerow(['id', 'label'])
for id in test_id_list:
feature = torch.tensor(test_json[id]['feature']).unsqueeze(0)
output = model(feature).squeeze()
prediction = 1 if torch.sigmoid(output) >= 0.5 else 0
csv_writer.writerow([id, prediction])
题目要求了95%的准确率,实测过程中训练损失可能会有些不稳定。总的来说,Torch代码的书写较keras还是繁琐了一些,后续还可以试试用keras重新回顾一下这段代码。