赞
踩
# -*- 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()
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!")
# -*- 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() ])
# -*- 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
# -*- 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
# -*- 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()
# -*- 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()
# -*- 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()
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。