这里提供一个baseline;
比赛地址:http://challenge.xfyun.cn/topic/info?type=PET
base = './train'
count = 0
with open(os.path.join('train_list.txt'), 'w') as f:
for i, cls_fold in enumerate(['AD', 'CN']):
try:
cls_base = os.path.join(base, cls_fold)
files = os.listdir(cls_base)
print('{} train num:'.format(cls_fold), len(files))
for pt in files:
img = os.path.join(cls_fold, pt)
info = img + ' ' + str(i) + '\n'
f.write(info)
count += 1
except Exception as e:
continue
print(count)
with open(os.path.join('labels.txt'), 'w') as f:
for i, cls_fold in enumerate(['AD', 'CN']):
f.write(cls_fold+'\n')
print('done.')
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.utils.data import DataLoader
import torchvision
from torchvision.transforms import transforms
from torchvision import models
from torchvision.models import ResNet
import numpy as np
import matplotlib.pyplot as plt
import os
data_dir = './'
train_dataset = torchvision.datasets.ImageFolder(root=os.path.join(data_dir, 'train'),
transform=transforms.Compose(
[
transforms.Resize(299),
transforms.ToTensor(),
transforms.Normalize(
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225))
]))
val_dataset = torchvision.datasets.ImageFolder(root=os.path.join(data_dir, 'train'),
transform=transforms.Compose(
[
transforms.Resize(299),
transforms.ToTensor(),
transforms.Normalize(
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225))
]))
train_dataloader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
val_dataloader = DataLoader(dataset=val_dataset, batch_size=1, shuffle=True)
# 类别名称
class_names = train_dataset.classes
print('class_names:{}'.format(class_names))
# 训练设备 CPU/GPU
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('trian_device:{}'.format(device.type))
这里使用vgg16:
# -------------------------模型选择,优化方法, 学习率策略----------------------
model_name = 'vgg16'
model = models.vgg16(pretrained=True)
num_classes = 2
# 改变全连接层,2分类问题,out_features = 2
# model.aux_logits = False
# model.AuxLogits.fc = nn.Linear(768, num_classes)
# model.fc = nn.Linear(2048, num_classes)
model.classifier[6] = nn.Linear(4096,num_classes)
# 模型迁移到CPU/GPU
model = model.to(device)
# 定义损失函数
loss_fc = nn.CrossEntropyLoss()
# 选择优化方法
optimizer = optim.Adam(model.parameters(), lr=2e-4)
# 学习率调整策略
exp_lr_scheduler = lr_scheduler.StepLR(optimizer=optimizer, step_size=30, gamma=0.5) # step_size
print('done.')
# ----------------训练过程-----------------
num_epochs = 101
for epoch in range(num_epochs):
running_loss = 0.0
exp_lr_scheduler.step()
for i, sample_batch in enumerate(train_dataloader):
inputs = sample_batch[0]
labels = sample_batch[1]
model.train()
# GPU/CPU
inputs = inputs.to(device)
labels = labels.to(device)
optimizer.zero_grad()
# foward
outputs = model(inputs)
# loss
loss = loss_fc(outputs, labels)
# loss求导,反向
loss.backward()
# 优化
optimizer.step()
#
running_loss += loss.item()
# 測試
if epoch % 5 == 0 and i == 1:
correct = 0
total = 0
model.eval()
for images_test, labels_test in val_dataloader:
images_test = images_test.to(device)
labels_test = labels_test.to(device)
outputs_test = model(images_test)
_, prediction = torch.max(outputs_test, 1)
correct += (torch.sum((prediction == labels_test))).item()
# print(prediction, labels_test, correct)
total += labels_test.size(0)
print('[{}, {}] running_loss = {:.5f} accurcay = {:.5f}'.format(epoch + 1, i + 1, running_loss / 20,
correct / total))
running_loss = 0.0
print('training finish !')
torch.save(model.state_dict(), './model.pth')
test_base = './test/AD&CN'
import pandas as pd
import cv2
from PIL import Image
from torchvision import transforms
from torch.autograd import Variable as V
import torch as t
import pickle as pkl
# model_name = 'vgg16_bn'
trans=transforms.Compose(
[
transforms.Resize(299),
transforms.ToTensor(),
transforms.Normalize(
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225))
]
)
model.eval()
submission = {'uuid':[], 'label':[]}
for i in tqdm(range(1, 1000+1)):
im_pt = os.path.join(test_base, '{}.png'.format(i))
img = cv2.imread(im_pt)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = Image.fromarray(img)
img = trans(img)
img = img.unsqueeze(0)
img = V(img.cuda())
predict = model(img)
predict = F.softmax(predict)[0]
temp = {'AD':float(predict[0]), 'CN':float(predict[1]), 'id':i}
_, pred = torch.max(predict, 0)
pred = class_names[int(pred)]
submission['uuid'].append(i)
submission['label'].append(pred)
data = pd.DataFrame(submission)
data.to_csv('./submission_{}.csv'.format(model_name), index=False)
print('done.')
大概 0.71 左右;