import argparse import numpy as np import os import pandas as pd from tqdm import tqdm import yaml import torch import torch.nn as nn from torch.utils.data import DataLoader from torch.utils.data.sampler import SequentialSampler, RandomSampler import time from sklearn.metrics import accuracy_score from contextlib import contextmanager from sklearn.model_selection import StratifiedKFold import matplotlib.pyplot as plt from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, CosineAnnealingLR, ReduceLROnPlateau from albumentations import ( HorizontalFlip, VerticalFlip, IAAPerspective, ShiftScaleRotate, CLAHE, RandomRotate90, Transpose, ShiftScaleRotate, Blur, OpticalDistortion, GridDistortion, HueSaturationValue, IAAAdditiveGaussianNoise, GaussNoise, MotionBlur, MedianBlur, IAAPiecewiseAffine, RandomResizedCrop, IAASharpen, IAAEmboss, RandomBrightnessContrast, Flip, OneOf, Compose, Normalize, Cutout, CoarseDropout, ShiftScaleRotate, CenterCrop, Resize ) from albumentations.pytorch import ToTensorV2 from albumentations import ImageOnlyTransform import math from torch.optim import Adam, SGD from utils import seed_everything import dataset import models parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument("--cfg", default='configs/resnext50_32x4d.yaml', type=str) parser.add_argument("--cuda", default='cuda:0', type=str) args = parser.parse_args() print(args.cuda) SEED = 123 seed_everything(SEED) if __name__ == "__main__": print(os.getcwd()) with open(args.cfg) as f: cfg = yaml.load(f, Loader=yaml.FullLoader) print(cfg) MODEL_DIR = 'weights/' + cfg['model_name'] + '_{}_{}_{}/'.format(cfg['size'], cfg['scheduler'], cfg['criterion']) OUTPUT_DIR = 'weights/' + cfg['model_name'] + '_{}_{}_{}/'.format(cfg['size'], cfg['scheduler'], cfg['criterion']) TRAIN_PATH = '/media/daitran/Data_SSD/data/SmartinsideAI/2021AIChamp/11.한국전력공사_데이터/train' TEST_PATH = '/media/daitran/Data_SSD/data/SmartinsideAI/2021AIChamp/11.한국전력공사_데이터/test' # ==================================================== # Utils # ==================================================== def get_score(y_true, y_pred): return accuracy_score(y_true, y_pred) @contextmanager def timer(name): t0 = time.time() LOGGER.info(f'[{name}] start') yield LOGGER.info(f'[{name}] done in {time.time() - t0:.0f} s.') def init_logger(log_file=OUTPUT_DIR+'inference.log'): from logging import getLogger, INFO, FileHandler, Formatter, StreamHandler logger = getLogger(__name__) logger.setLevel(INFO) handler1 = StreamHandler() handler1.setFormatter(Formatter("%(message)s")) handler2 = FileHandler(filename=log_file) handler2.setFormatter(Formatter("%(message)s")) logger.addHandler(handler1) logger.addHandler(handler2) return logger test = pd.read_csv('/media/daitran/Data_SSD/data/SmartinsideAI/2021AIChamp/11.한국전력공사_데이터/test_output_sample.csv') test_dataset = dataset.TestDataset(test, transform=None) def get_transforms(*, data): if data == 'train': return Compose([ Resize(cfg['size'], cfg['size']), # RandomResizedCrop(cfg['size'], cfg['size']), Transpose(p=0.5), HorizontalFlip(p=0.5), VerticalFlip(p=0.5), ShiftScaleRotate(p=0.5), Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], ), ToTensorV2(), ]) elif data == 'valid': return Compose([ Resize(cfg['size'], cfg['size']), CLAHE(), HorizontalFlip(p=0.5), VerticalFlip(p=0.5), Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], ), ToTensorV2(), ]) def inference(model, states, test_loader, device): model.to(device) tk0 = tqdm(enumerate(test_loader), total=len(test_loader)) probs = [] for i, (images) in tk0: images = images.to(device) avg_preds = [] for state in states: model.load_state_dict(state['model']) model.eval() with torch.no_grad(): y_preds = model(images) avg_preds.append(y_preds.softmax(1).to('cpu').numpy()) avg_preds = np.mean(avg_preds, axis=0) probs.append(avg_preds) probs = np.concatenate(probs) return probs device = 'cuda:0' # ==================================================== # model & optimizer # ==================================================== # ResNext if cfg['model_name'][:7] == 'resnext': model = models.CustomResNext(cfg['model_name'], pretrained=True) # ViT elif cfg['model_name'][:3] == 'vit': model = models.VITModel(num_classes=cfg['target_size'], model_name=cfg['model_name']) # EfficientNet elif cfg['model_name'][:12] == 'efficientnet' or cfg['model_name'][3:12+3] == 'efficientnet': model = models.EfficientNetModel(num_classes=cfg['target_size'], model_name=cfg['model_name']) # Seresnext elif cfg['model_name'][:9] == 'seresnext': model = models.SeResNext(num_classes=cfg['target_size'], model_name=cfg['model_name']) model.to(device) #################################################### m_name = cfg['model_name'] # Update map location states = [torch.load(MODEL_DIR+f'{m_name}_fold{fold}_best.pth', map_location='cuda:0') for fold in cfg['trn_fold']] # Customize -> loadhigh kfold only # states = [torch.load('/media/daitran/Data_SSD/code/SmartInsideAI/github/2021aichamp/src/classify_status_models/weights/tf_efficientnetv2_l_in21ft1k_512_CosineAnnealingWarmRestarts_TaylorCrossEntropyLoss/tf_efficientnetv2_l_in21ft1k_fold0_best.pth')] test_dataset = dataset.TestDataset(test, transform=get_transforms(data='valid')) test_loader = DataLoader(test_dataset, batch_size=cfg['batch_size'], shuffle=False, num_workers=cfg['num_workers'], pin_memory=True) print("Inferencing") device = 'cuda:0' print(os.getcwd()) tta_predictions = [] for t in range(cfg['tta']): print("TTA {}".format(t)) predictions = inference(model, states, test_loader, device) tta_predictions.append(predictions) final_preds = np.mean(tta_predictions, axis = 0) # print(predictions.argmax(1)) test['status'] = final_preds.argmax(1) submission = test.drop(["fault", "type"], axis = 1) OUTPUT_DIR = 'weights/' + cfg['model_name'] + '_{}_{}_{}/'.format(cfg['size'], cfg['scheduler'], cfg['criterion']) submission.to_csv(OUTPUT_DIR + "/submission_status.csv", index= False) # save final preds probs print("Save at: {}".format(OUTPUT_DIR)) final_preds_csv = pd.DataFrame(final_preds) final_preds_csv.to_csv(OUTPUT_DIR + "/final_preds.csv", index = False)