target_small_value, target_small_class = torch.max(
target_labels[big_class_list[i]][i].view([1, -1]), dim=1)
res = (pred_small_class == target_small_class).data[0]
if res:
true_num += 1
return true_num / batch_size
if __name__ == "__main__":
tran = transforms.Compose(
[transforms.ToTensor(),
transforms.Resize((224, 224))])
reader = ImageReader.MultiTaskImageReader(
root_path="../../ali_data/second_validate_data/Images/",
csv_path="../../ali_data/second_validate_data/labels.csv",
transform=tran,
name_pos_dict=config.name_pos_dict,
name_num_dict=config.name_num_dict,
num_list=config.num_list)
trainloader = torch.utils.data.DataLoader(reader,
batch_size=1,
shuffle=True,
num_workers=4)
for x, y in trainloader:
break
model = Network.MyDenseNet(config.name_num_dict)
x = torch.autograd.Variable(x)
output = model(x)
for i in range(len(output)):
y[i] = torch.autograd.Variable(y[i])
def train():
'''
定义训练的函数
'''
#加载模型
model = Network.MyDenseNet(name_num_dict=config.name_num_dict,
cnn_weights_path=config.cnn_weights_path)
#加载全部的模型权重
if config.all_weights_path:
weights_dict = torch.load(config.all_weights_path)
model.load_state_dict(state_dict=weights_dict)
print("load {} weights succecced...".format(config.all_weights_path))
if torch.cuda.is_available() and use_gpu:
model = model.cuda()
#定义训练数据部分
tran = transforms.Compose(
[transforms.ToTensor(),
transforms.Resize(config.target_size)])
reader = ImageReader.MultiTaskImageReader(
root_path=config.train_img_root_path,
csv_path=config.train_img_csv,
transform=tran,
name_pos_dict=config.name_pos_dict,
name_num_dict=config.name_num_dict,
num_list=config.num_list)
trainloader = torch.utils.data.DataLoader(reader,
batch_size=config.batch_size,
shuffle=True,
num_workers=config.num_thread)
#定义验证数据部分
val_reader = ImageReader.MultiTaskImageReader(
root_path=config.val_img_root_path,
csv_path=config.val_img_csv,
transform=tran,
name_pos_dict=config.name_pos_dict,
name_num_dict=config.name_num_dict,
num_list=config.num_list)
validateloader = torch.utils.data.DataLoader(val_reader,
batch_size=config.batch_size,
shuffle=True,
num_workers=config.num_thread)
#定义优化器
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
for epoch in range(config.start_epoch, config.start_epoch + config.epochs):
print("{} epoch start...".format(epoch))
#开始训练模式
model.train()
#定义一轮中的评估指标
step = 0 #训练步数
total_loss = 0 #记录总的运行损失
total_loss_list = [0, 0, 0, 0, 0, 0, 0, 0] #记录各个大类的运行损失
total_acc = 0 #记录总的准确率
val_step = 0
val_total_acc = 0
val_total_loss = 0
val_total_loss_list = [0, 0, 0, 0, 0, 0, 0, 0]
for data in trainloader:
inputs, labels = data
inputs = Variable(inputs)
for i in range(len(labels)):
labels[i] = Variable(labels[i])
if torch.cuda.is_available() and use_gpu:
inputs = inputs.cuda()
for i in range(len(labels)):
labels[i] = labels[i].cuda()
optimizer.zero_grad() #将梯度缓冲区清0
outputs = model(inputs)
step_loss, step_loss_list = metrics.loss(outputs, labels)
total_loss += step_loss.data[0] #更新总的损失
for i in range(len(total_loss_list)): #更新各个类的损失
total_loss_list[i] += step_loss_list[i]
step_loss.backward()
optimizer.step()
if step % 10 == 0 and step != 0:
step_acc = metrics.accuracy(outputs, labels)
total_acc += step_acc
print("{} epoch, {} step, step loss is {}, step acc is {}".
format(epoch, step, step_loss.data[0], step_acc))
else:
print("{} epoch, {} step, step loss is {}".format(
epoch, step, step_loss.data[0]))
#释放显存
del ([inputs, labels, outputs, step_loss, step_loss_list])
step += 1
#模型保存参数
print("start save model weights...")
torch.save(model.state_dict(), config.save_weights_path.format(epoch))
model.eval() #开启验证模式
print("start validate ...")
#解下来进行验证
for data in validateloader:
inputs, labels = data
inputs = Variable(inputs)
for i in range(len(labels)):
labels[i] = Variable(labels[i])
if torch.cuda.is_available() and use_gpu:
inputs = inputs.cuda()
for i in range(len(labels)):
labels[i] = labels[i].cuda()
#获取outputs
outputs = model(inputs)
step_loss, step_loss_list = metrics.loss(outputs, labels)
step_acc = metrics.accuracy(outputs, labels)
val_total_loss += step_loss.data[0]
for i in range(len(val_total_loss_list)):
val_total_loss_list[i] += step_loss_list[i]
val_total_acc += step_acc
val_step += 1
#释放显存
del ([inputs, labels, outputs, step_loss, step_loss_list])
print("{} epoch , validation loss is {}, validation accuracy is {}".
format(epoch, val_total_loss / val_step,
val_total_acc / val_step))
#训练完一轮之后将训练结果记录到文件中
print("start log train info...")
utils.log_train(epoch, total_loss, total_loss_list, step, total_acc,
val_total_loss, val_total_loss_list, val_total_acc,
val_step)
del (model)
print("{} epoch train end ...".format(epoch))