参数设置，config.py
# -*- coding: utf-8 -*-
# @Time : 2021/1/20 19:07
# @Author : Johnson
class DefaultConfigs(object):
#1.string parameters
train_data = "/home/user/zls/data/train/"
test_data = ""
val_data = "/home/user/zls/data/val/"
model_name = "resnet50"
weights = "./checkpoints/"
best_models = weights + "best_model/"
submit = "./submit/"
logs = "./logs/"
gpus = "1"
augmen_level = "medium" # "light","hard","hard2"
#2.numeric parameters
epochs = 40
batch_size = 16
img_height = 300
img_weight = 300
num_classes = 62
seed = 888
lr = 1e-4
lr_decay = 1e-4
weight_decay = 1e-4
config = DefaultConfigs()
数据加载：data_loader.py
pytorch 的数据读取方式有两种，一种是不同类别的图像按照文件夹进行划分，比如交通标志数据集：
train/
00000/
01153_00000.png
01153_00001.png
00001/
00025_00000.png
00025_00001.png
# -*- coding: utf-8 -*-
# @Time : 2021/1/20 19:01
# @Author : Johnson
from torch.utils.data import Dataset
from torchvision import transforms as T
from config import config
from PIL import Image
from itertools import chain
from glob import glob
from tqdm import tqdm
from .augmentations import get_train_transform,get_test_transform
import random
import numpy as np
import pandas as pd
import os
import cv2
import torch
#1.set random seed
random.seed(config.seed)
np.random.seed(config.seed)
torch.manual_seed(config.seed)
torch.cuda.manual_seed_all(config.seed)
#2.define dataset
class ChaojieDataset(Dataset):
def __init__(self,label_list,train=True,test=False):
self.test = test
self.train = train
imgs = []
if self.test:
for index,row in label_list.iterrows():
imgs.append(row["filename"])
self.imgs = imgs
else:
for index,row in label_list.iterrows():
imgs.append(row['filename',row["label"]])
self.imgs = imgs
def __getitem__(self,index):
if self.test:
filename = self.imgs[index]
img = cv2.imread(filename)
img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
img = cv2.resize(img,(int(config.img_height*1.5),int(config.img_weight*1.5)))
img = get_test_transform(img.shape)(image=img)["image"]
return img,filename
else:
filename,label = self.imgs[index]
img = cv2.imread(filename)
img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
img = cv2.resize(img, (int(config.img_height * 1.5), int(config.img_weight * 1.5)))
img = get_train_transform(img.shape, augmentation=config.augmen_level)(image=img)["image"]
return img, label
def __len__(self):
return len(self.imgs)
def collate_fn(batch):
imgs = []
label = []
for sample in batch:
imgs.append(sample[0])
label.append(sample[1])
return torch.stack(imgs,0),label
def get_files(root,mode):
#for test
if mode == "test":
files = []
for img in os.listdir(root):
files.append(root + img)
files = pd.DataFrame({"filename":files})
return files
elif mode != "test":
#for train and val
all_data_path,labels = [],[]
image_folders = list(map(lambda x:root+x,os.listdir(root)))
all_images = list(chain.from_iterable(list(map(lambda x:glob(x+"/*"),image_folders))))
print("loading train dataset")
for file in tqdm(all_images):
all_data_path.append(file)
labels.append(int(file.split("/")[-2]))
all_files = pd.DataFrame({"filename":all_data_path,"label":labels})
return all_files
else:
print("check the mode please!")
数据增强：augmentations.py
# -*- coding: utf-8 -*-
# @Time : 2021/1/20 19:00
# @Author : Johnson
import random
from typing import List,Tuple
import albumentations as A
import cv2
import numpy as np
from albumentations.augmentations.functional import brightness_contrast_adjust,elastic_transform
from albumentations.pytorch import ToTensor
from albumentations.augmentations.transforms import Resize,CenterCrop
from config import config
class IndependentRandomBrightnessContrast(A.ImageOnlyTransform):
""" Change brightness & contrast independently per channels """
def __init__(self, brightness_limit=0.2, contrast_limit=0.2, always_apply=False, p=0.5):
super(IndependentRandomBrightnessContrast, self).__init__(always_apply, p)
self.brightness_limit = A.to_tuple(brightness_limit)
self.contrast_limit = A.to_tuple(contrast_limit)
def apply(self, img, **params):
img = img.copy()
for ch in range(img.shape[2]):
alpha = 1.0 + random.uniform(self.contrast_limit[0], self.contrast_limit[1])
beta = 0.0 + random.uniform(self.brightness_limit[0], self.brightness_limit[1])
img[..., ch] = brightness_contrast_adjust(img[..., ch], alpha, beta)
return img
def get_light_augmentations(image_size):
return A.Compose([
A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.1,
rotate_limit=15,
border_mode=cv2.BORDER_CONSTANT, value=0),
A.RandomSizedCrop(min_max_height=(int(image_size[0] * 0.85), image_size[0]),
height=image_size[0],
width=image_size[1], p=0.3),
# Brightness/contrast augmentations
A.OneOf([
A.RandomBrightnessContrast(brightness_limit=0.25,
contrast_limit=0.2),
IndependentRandomBrightnessContrast(brightness_limit=0.1,
contrast_limit=0.1),
A.RandomGamma(gamma_limit=(75, 125)),
A.NoOp()
]),
A.HorizontalFlip(p=0.5),
])
def get_medium_augmentations(image_size):
return A.Compose([
A.OneOf([
A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.1,
rotate_limit=15,
border_mode=cv2.BORDER_CONSTANT, value=0),
A.OpticalDistortion(distort_limit=0.11, shift_limit=0.15,
border_mode=cv2.BORDER_CONSTANT,
value=0),
A.NoOp()
]),
A.RandomSizedCrop(min_max_height=(int(image_size[0] * 0.75), image_size[0]),
height=image_size[0],
width=image_size[1], p=0.3),
A.OneOf([
A.RandomBrightnessContrast(brightness_limit=0.5,
contrast_limit=0.4),
IndependentRandomBrightnessContrast(brightness_limit=0.25,
contrast_limit=0.24),
A.RandomGamma(gamma_limit=(50, 150)),
A.NoOp()
]),
A.OneOf([
A.RGBShift(r_shift_limit=20, b_shift_limit=15, g_shift_limit=15),
A.HueSaturationValue(hue_shift_limit=5,
sat_shift_limit=5),
A.NoOp()
]),
A.HorizontalFlip(p=0.5),
A.VerticalFlip(p=0.5)
])
def get_hard_augmentations(image_size):
return A.Compose([
A.OneOf([
A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.1,
rotate_limit=45,
border_mode=cv2.BORDER_CONSTANT, value=0),
A.ElasticTransform(alpha_affine=0,
alpha=35,
sigma=5,
border_mode=cv2.BORDER_CONSTANT,
value=0),
A.OpticalDistortion(distort_limit=0.11, shift_limit=0.15,
border_mode=cv2.BORDER_CONSTANT,
value=0),
A.GridDistortion(border_mode=cv2.BORDER_CONSTANT,
value=0),
A.NoOp()
]),
A.OneOf([
A.RandomSizedCrop(min_max_height=(int(image_size[0] * 0.75), image_size[0]),
height=image_size[0],
width=image_size[1], p=0.3),
A.NoOp()
]),
A.ISONoise(p=0.5),
# Brightness/contrast augmentations
A.OneOf([
A.RandomBrightnessContrast(brightness_limit=0.5,
contrast_limit=0.4),
IndependentRandomBrightnessContrast(brightness_limit=0.25,
contrast_limit=0.24),
A.RandomGamma(gamma_limit=(50, 150)),
A.NoOp()
]),
A.OneOf([
A.RGBShift(r_shift_limit=40, b_shift_limit=30, g_shift_limit=30),
A.HueSaturationValue(hue_shift_limit=10,
sat_shift_limit=10),
A.ToGray(p=0.2),
A.NoOp()
]),
A.ChannelDropout(),
A.RandomGridShuffle(p=0.3),
# D4
A.Compose([
A.RandomRotate90(),
A.Transpose()
])
])
def get_hard_augmentations_v2(image_size):
return A.Compose([
A.OneOf([
A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.1,
rotate_limit=45,
border_mode=cv2.BORDER_CONSTANT, value=0),
A.ElasticTransform(alpha_affine=0,
alpha=35,
sigma=5,
border_mode=cv2.BORDER_CONSTANT,
value=0),
A.OpticalDistortion(distort_limit=0.11, shift_limit=0.15,
border_mode=cv2.BORDER_CONSTANT,
value=0),
A.GridDistortion(border_mode=cv2.BORDER_CONSTANT,
value=0),
A.NoOp()
]),
A.OneOf([
A.RandomSizedCrop(min_max_height=(int(image_size[0] * 0.75), image_size[0]),
height=image_size[0],
width=image_size[1], p=0.3),
A.NoOp()
]),
A.ISONoise(p=0.5),
A.JpegCompression(p=0.3, quality_lower=75),
# Brightness/contrast augmentations
A.OneOf([
A.RandomBrightnessContrast(brightness_limit=0.5,
contrast_limit=0.4),
IndependentRandomBrightnessContrast(brightness_limit=0.25,
contrast_limit=0.24),
A.RandomGamma(gamma_limit=(50, 150)),
A.NoOp()
]),
A.OneOf([
A.RGBShift(r_shift_limit=40, b_shift_limit=30, g_shift_limit=30),
A.HueSaturationValue(hue_shift_limit=10,
sat_shift_limit=10),
A.ToGray(p=0.2),
A.NoOp()
]),
A.OneOf([
A.ChannelDropout(p=0.2),
A.CoarseDropout(p=0.1, max_holes=2, max_width=256, max_height=256, min_height=16, min_width=16),
A.NoOp()
]),
A.RandomGridShuffle(p=0.3),
# D4
A.Compose([
A.RandomRotate90(),
A.Transpose()
])
])
def get_none_augmentations(image_size):
return A.NoOp()
def get_train_transform(image_size,augmentation=None):
if augmentation is None:
augmentation = 'none'
LEVELS = {
'none':get_none_augmentations,
'light':get_light_augmentations,
'medium':get_medium_augmentations,
'hard':get_hard_augmentations,
'hard2':get_hard_augmentations_v2
}
assert augmentation is LEVELS.keys()
augmentation = LEVELS[augmentation](image_size)
longest_size = max(image_size[0],image_size[1])
return A.Compose([
# Resize(int(config.img_height*1.5),int(config.img_weight*1.5)),
CenterCrop(config.img_height, config.img_weight),
A.LongestMaxSize(longest_size, interpolation=cv2.INTER_CUBIC),
A.PadIfNeeded(image_size[0], image_size[1],
border_mode=cv2.BORDER_CONSTANT, value=0),
augmentation,
A.Normalize(),
ToTensor()
])
def get_test_transform(image_size):
longest_size = max(image_size[0], image_size[1])
return A.Compose([
#Resize(int(config.img_height*1.5),int(config.img_weight*1.5)),
CenterCrop(config.img_height,config.img_weight),
A.LongestMaxSize(longest_size, interpolation=cv2.INTER_CUBIC),
A.PadIfNeeded(image_size[0], image_size[1],
border_mode=cv2.BORDER_CONSTANT, value=0),
A.Normalize(),
ToTensor()
])
数据加载:model.py
# -*- coding: utf-8 -*-
# @Time : 2021/1/20 20:55
# @Author : Johnson
import torchvision
import torch.nn.functional as F
from torch import nn
from config import config
def generate_model():
class DenseModel(nn.Module):
def __init__(self, pretrained_model):
super(DenseModel, self).__init__()
self.classifier = nn.Linear(pretrained_model.classifier.in_features, config.num_classes)
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal(m.weight)
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
elif isinstance(m, nn.Linear):
m.bias.data.zero_()
self.features = pretrained_model.features
self.layer1 = pretrained_model.features._modules['denseblock1']
self.layer2 = pretrained_model.features._modules['denseblock2']
self.layer3 = pretrained_model.features._modules['denseblock3']
self.layer4 = pretrained_model.features._modules['denseblock4']
def forward(self,x):
features = self.features(x)
out = F.relu(features,inplace=True)
out = F.avg_pool2d(out,kernel_size=8).view(features.size(0),-1)
out = F.sigmoid(self.classifier(out))
return out
return DenseModel(torchvision.models.densenet169(pretrained=True))
def get_net():
#return MyModel(torchvision.models.resnet101(pretrained = True))
model = torchvision.models.resnet152(pretrained = True)
#for param in model.parameters():
# param.requires_grad = False
model.avgpool = nn.AdaptiveAvgPool2d(1)
model.fc = nn.Linear(2048,config.num_classes)
return model
指标评价：utils.py
# -*- coding: utf-8 -*-
# @Time : 2021/1/20 21:04
# @Author : Johnson
import shutil
import torch
import sys
import os
from config import config
def save_checkpoint(state, is_best,fold):
filename = config.weights + config.model_name + os.sep +str(fold) + os.sep + "_checkpoint.pth.tar"
torch.save(state, filename)
if is_best:
message = config.best_models + config.model_name+ os.sep +str(fold) + os.sep + 'model_best.pth.tar'
print("Get Better top1 : %s saving weights to %s"%(state["best_precision1"],message))
with open("./logs/%s.txt"%config.model_name,"a") as f:
print("Get Better top1 : %s saving weights to %s"%(state["best_precision1"],message),file=f)
shutil.copyfile(filename, message)
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
def adjust_learning_rate(optimizer,epoch):
"""sets the learning rate to the initial LR decayed by 10 every 3 epochs """
lr = config.lr*(0.1**(epoch//3))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
def schedule(current_epoch, current_lrs, **logs):
lrs = [1e-3, 1e-4, 0.5e-4, 1e-5, 0.5e-5]
epochs = [0, 1, 6, 8, 12]
for lr, epoch in zip(lrs, epochs):
if current_epoch >= epoch:
current_lrs[5] = lr
if current_epoch >= 2:
current_lrs[4] = lr * 1
current_lrs[3] = lr * 1
current_lrs[2] = lr * 1
current_lrs[1] = lr * 1
current_lrs[0] = lr * 0.1
return current_lrs
def accuracy(output, target, topk=(1,)):
"""Computes the accuracy over the k top predictions for the specified values of k"""
with torch.no_grad():
maxk = max(topk)
batch_size = target.size(0)
_, pred = output.topk(maxk, 1, True, True)
pred = pred.t()
correct = pred.eq(target.view(1, -1).expand_as(pred))
res = []
for k in topk:
correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
res.append(correct_k.mul_(100.0 / batch_size))
return res
def get_learning_rate(optimizer):
lr=[]
for param_group in optimizer.param_groups:
lr +=[ param_group['lr'] ]
#assert(len(lr)==1) #we support only one param_group
lr = lr[0]
return lr
def time_to_str(t, mode='min'):
if mode=='min':
t = int(t)/60
hr = t//60
min = t%60
return '%2d hr %02d min'%(hr,min)
elif mode=='sec':
t = int(t)
min = t//60
sec = t%60
return '%2d min %02d sec'%(min,sec)
else:
raise NotImplementedError
进度条:progress.py
# -*- coding: utf-8 -*-
# @Time : 2021/1/21 9:59
# @Author : Johnson
import sys
import re
class ProgressBar(object):
DEFAULT = "Progress: %(bar)s %(percent)3d%%"
def __init__(self, mode, epoch=None, total_epoch=None, current_loss=None, current_top1=None, model_name=None,
total=None, current=None, width=50, symbol=">", output=sys.stderr):
assert len(symbol) == 1
self.mode = mode
self.total = total
self.symbol = symbol
self.output = output
self.width = width
self.current = current
self.epoch = epoch
self.total_epoch = total_epoch
self.current_loss = current_loss
self.current_top1 = current_top1
self.model_name = model_name
def __call__(self):
percent = self.current / float(self.total)
size = int(self.width * percent)
bar = "[" + self.symbol * size + " " * (self.width - size) + "]"
args = {
"mode": self.mode,
"total": self.total,
"bar": bar,
"current": self.current,
"percent": percent * 100,
"current_loss": self.current_loss,
"current_top1": self.current_top1,
"epoch": self.epoch + 1,
"epochs": self.total_epoch
}
message = "\033[1;32;40m%(mode)s Epoch: %(epoch)d/%(epochs)d %(bar)s\033[0m [Current: Loss %(current_loss)f Top1: %(current_top1)f ] %(current)d/%(total)d \033[1;32;40m[ %(percent)3d%% ]\033[0m" % args
self.write_message = "%(mode)s Epoch: %(epoch)d/%(epochs)d %(bar)s [Current: Loss %(current_loss)f Top1: %(current_top1)f ] %(current)d/%(total)d [ %(percent)3d%% ]" % args
print("\r" + message, file=self.output, end="")
def done(self):
self.current = self.total
self()
print("", file=self.output)
with open("./logs/%s.txt" % self.model_name, "a") as f:
print(self.write_message, file=f)
if __name__ == "__main__":
from time import sleep
progress = ProgressBar("Train", total_epoch=150, model_name="resnet159")
for i in range(150):
progress.total = 50
progress.epoch = i
progress.current_loss = 0.15
progress.current_top1 = 0.45
for x in range(50):
progress.current = x
progress()
sleep(0.1)
progress.done()
主函数:main.py
# -*- coding: utf-8 -*-
# @Time : 2021/1/21 10:02
# @Author : Johnson
import os
import random
import time
import json
import torch
import torchvision
import numpy as np
import pandas as pd
import warnings
from datetime import datetime
from torch import nn,optim
from config import config
from collections import OrderedDict
from torch.autograd import Variable
from torch.utils.data import DataLoader
from dataset.dataloader import *
from sklearn.model_selection import train_test_split,StratifiedKFold
from timeit import default_timer as timer
from models.model import *
from utils import *
from IPython import embed
#1. set random.seed and cudnn performance
random.seed(config.seed)
np.random.seed(config.seed)
torch.manual_seed(config.seed)
torch.cuda.manual_seed_all(config.seed)
os.environ["CUDA_VISIBLE_DEVICES"] = config.gpus
torch.backends.cudnn.benchmark = True
warnings.filterwarnings('ignore')
# 2. evaluate func
def evaluate(val_loader, model, criterion, epoch):
# 2.1 define meters
losses = AverageMeter()
top1 = AverageMeter()
# progress bar
val_progressor = ProgressBar(mode="Val ", epoch=epoch, total_epoch=config.epochs, model_name=config.model_name,
total=len(val_loader))
# 2.2 switch to evaluate mode and confirm model has been transfered to cuda
model.cuda()
model.eval()
with torch.no_grad():
for i, (input, target) in enumerate(val_loader):
val_progressor.current = i
input = Variable(input).cuda()
target = Variable(torch.from_numpy(np.array(target)).long()).cuda()
# target = Variable(target).cuda()
# 2.2.1 compute output
output = model(input)
loss = criterion(output, target)
# 2.2.2 measure accuracy and record loss
precision1, precision2 = accuracy(output, target, topk=(1, 2))
losses.update(loss.item(), input.size(0))
top1.update(precision1[0], input.size(0))
val_progressor.current_loss = losses.avg
val_progressor.current_top1 = top1.avg
val_progressor()
val_progressor.done()
return [losses.avg, top1.avg]
# 3. test model on public dataset and save the probability matrix
def test(test_loader, model, folds):
# 3.1 confirm the model converted to cuda
csv_map = OrderedDict({"filename": [], "probability": []})
model.cuda()
model.eval()
for i, (input, filepath) in enumerate(tqdm(test_loader)):
# 3.2 change everything to cuda and get only basename
filepath = [os.path.basename(x) for x in filepath]
with torch.no_grad():
image_var = Variable(input).cuda()
# 3.3.output
# print(filepath)
# print(input,input.shape)
y_pred = model(image_var)
print(y_pred.shape)
smax = nn.Softmax(1)
smax_out = smax(y_pred)
# 3.4 save probability to csv files
csv_map["filename"].extend(filepath)
for output in smax_out:
prob = ";".join([str(i) for i in output.data.tolist()])
csv_map["probability"].append(prob)
result = pd.DataFrame(csv_map)
result["probability"] = result["probability"].map(lambda x: [float(i) for i in x.split(";")])
result.to_csv("./submit/{}_submission.csv".format(config.model_name + "_" + str(folds)), index=False, header=None)
# 4. more details to build main function
def main():
fold = 0
# 4.1 mkdirs
if not os.path.exists(config.submit):
os.mkdir(config.submit)
if not os.path.exists(config.weights):
os.mkdir(config.weights)
if not os.path.exists(config.best_models):
os.mkdir(config.best_models)
if not os.path.exists(config.logs):
os.mkdir(config.logs)
if not os.path.exists(config.weights + config.model_name + os.sep + str(fold) + os.sep):
os.makedirs(config.weights + config.model_name + os.sep + str(fold) + os.sep)
if not os.path.exists(config.best_models + config.model_name + os.sep + str(fold) + os.sep):
os.makedirs(config.best_models + config.model_name + os.sep + str(fold) + os.sep)
# 4.2 get model and optimizer
model = get_net()
# model = torch.nn.DataParallel(model)
model.cuda()
# optimizer = optim.SGD(model.parameters(),lr = config.lr,momentum=0.9,weight_decay=config.weight_decay)
optimizer = optim.Adam(model.parameters(), lr=config.lr, amsgrad=True, weight_decay=config.weight_decay)
criterion = nn.CrossEntropyLoss().cuda()
# 4.3 some parameters for K-fold and restart model
start_epoch = 0
best_precision1 = 0
best_precision_save = 0
resume = False
# 4.4 restart the training process
if resume:
checkpoint = torch.load(config.best_models + str(fold) + "/model_best.pth.tar")
start_epoch = checkpoint["epoch"]
fold = checkpoint["fold"]
best_precision1 = checkpoint["best_precision1"]
model.load_state_dict(checkpoint["state_dict"])
optimizer.load_state_dict(checkpoint["optimizer"])
# 4.5 get files and split for K-fold dataset
# 4.5.1 read files
train_data_list = get_files(config.train_data, "train")
val_data_list = get_files(config.val_data, "val")
# test_files = get_files(config.test_data,"test")
"""
#4.5.2 split
split_fold = StratifiedKFold(n_splits=3)
folds_indexes = split_fold.split(X=origin_files["filename"],y=origin_files["label"])
folds_indexes = np.array(list(folds_indexes))
fold_index = folds_indexes[fold]
#4.5.3 using fold index to split for train data and val data
train_data_list = pd.concat([origin_files["filename"][fold_index[0]],origin_files["label"][fold_index[0]]],axis=1)
val_data_list = pd.concat([origin_files["filename"][fold_index[1]],origin_files["label"][fold_index[1]]],axis=1)
"""
# train_data_list,val_data_list = train_test_split(origin_files,test_size = 0.1,stratify=origin_files["label"])
# 4.5.4 load dataset
train_dataloader = DataLoader(ChaojieDataset(train_data_list), batch_size=config.batch_size, shuffle=True,
collate_fn=collate_fn, pin_memory=True, num_workers=4)
val_dataloader = DataLoader(ChaojieDataset(val_data_list, train=False), batch_size=config.batch_size * 2,
shuffle=True, collate_fn=collate_fn, pin_memory=False, num_workers=4)
# test_dataloader = DataLoader(ChaojieDataset(test_files,test=True),batch_size=1,shuffle=False,pin_memory=False)
# scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,"max",verbose=1,patience=3)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
# 4.5.5.1 define metrics
train_losses = AverageMeter()
train_top1 = AverageMeter()
valid_loss = [np.inf, 0, 0]
model.train()
# 4.5.5 train
start = timer()
for epoch in range(start_epoch, config.epochs):
scheduler.step(epoch)
train_progressor = ProgressBar(mode="Train", epoch=epoch, total_epoch=config.epochs,
model_name=config.model_name, total=len(train_dataloader))
# train
# global iter
for iter, (input, target) in enumerate(train_dataloader):
# 4.5.5 switch to continue train process
train_progressor.current = iter
model.train()
input = Variable(input).cuda()
target = Variable(torch.from_numpy(np.array(target)).long()).cuda()
# target = Variable(target).cuda()
output = model(input)
loss = criterion(output, target)
precision1_train, precision2_train = accuracy(output, target, topk=(1, 2))
train_losses.update(loss.item(), input.size(0))
train_top1.update(precision1_train[0], input.size(0))
train_progressor.current_loss = train_losses.avg
train_progressor.current_top1 = train_top1.avg
# backward
optimizer.zero_grad()
loss.backward()
optimizer.step()
train_progressor()
train_progressor.done()
# evaluate
lr = get_learning_rate(optimizer)
# evaluate every half epoch
valid_loss = evaluate(val_dataloader, model, criterion, epoch)
is_best = valid_loss[1] > best_precision1
best_precision1 = max(valid_loss[1], best_precision1)
try:
best_precision_save = best_precision1.cpu().data.numpy()
except:
pass
save_checkpoint({
"epoch": epoch + 1,
"model_name": config.model_name,
"state_dict": model.state_dict(),
"best_precision1": best_precision1,
"optimizer": optimizer.state_dict(),
"fold": fold,
"valid_loss": valid_loss,
}, is_best, fold)
if __name__ == "__main__":
main()
测试函数:test.py
# -*- coding: utf-8 -*-
# @Time : 2021/1/21 10:06
# @Author : Johnson
import os
import random
import time
import json
import torch
import torchvision
import numpy as np
import pandas as pd
import warnings
from datetime import datetime
from torch import nn,optim
from config import config
from collections import OrderedDict
from torch.autograd import Variable
from torch.utils.data import DataLoader
from dataset.dataloader import *
from sklearn.model_selection import train_test_split,StratifiedKFold
from timeit import default_timer as timer
from models.model import *
from utils import *
from IPython import embed
#1. set random.seed and cudnn performance
random.seed(config.seed)
np.random.seed(config.seed)
torch.manual_seed(config.seed)
torch.cuda.manual_seed_all(config.seed)
os.environ["CUDA_VISIBLE_DEVICES"] = config.gpus
torch.backends.cudnn.benchmark = True
warnings.filterwarnings('ignore')
#3. test model on public dataset and save the probability matrix
def test(test_loader,model,folds):
#3.1 confirm the model converted to cuda
test_labels = open("%s.csv"%config.model_name,"w")
model.cuda()
model.eval()
tta = False
for i,(input,filepath) in enumerate(tqdm(test_loader)):
#3.2 change everything to cuda and get only basename
filepath = [os.path.basename(x) for x in filepath]
with torch.no_grad():
image_var = Variable(input).cuda()
#3.3.output
#print(filepath)
#print(input,input.shape)
if tta == False:
y_pred = model(image_var)
smax = nn.Softmax(1)
smax_out = smax(y_pred)
label = np.argmax(smax_out.cpu().data.numpy())
test_labels.write(filepath[0]+","+str(label)+"\n")
#4. more details to build main function
def main():
fold = 0
#4.1 mkdirs
if not os.path.exists(config.submit):
os.mkdir(config.submit)
if not os.path.exists(config.weights):
os.mkdir(config.weights)
if not os.path.exists(config.best_models):
os.mkdir(config.best_models)
if not os.path.exists(config.logs):
os.mkdir(config.logs)
if not os.path.exists(config.weights + config.model_name + os.sep +str(fold) + os.sep):
os.makedirs(config.weights + config.model_name + os.sep +str(fold) + os.sep)
if not os.path.exists(config.best_models + config.model_name + os.sep +str(fold) + os.sep):
os.makedirs(config.best_models + config.model_name + os.sep +str(fold) + os.sep)
#4.2 get model and optimizer
model = get_net()
#model = torch.nn.DataParallel(model)
model.cuda()
#4.5 get files and split for K-fold dataset
test_files = get_files(config.test_data,"test")
#4.5.4 load dataset
test_dataloader = DataLoader(ChaojieDataset(test_files,test=True),batch_size=1,shuffle=False,pin_memory=False)
best_model = torch.load("checkpoints/best_model/%s/0/model_best.pth.tar"%config.model_name)
model.load_state_dict(best_model["state_dict"])
test(test_dataloader,model,fold)
if __name__ =="__main__":
main()