# -*- coding:UTF-8 -*-
"""
@create: 2021-10-06 10:50:00
@author: GongChen
@description: 训练脚本
"""

import os
import torch
import torch.nn as nn
import torch.optim as optim
from torchsummary import summary

from models.vgg import *
from utils.arg_utils import *
from utils.data_utils import *
from utils.progress_utils import progress_bar
from utils.earlystopping import EarlyStopping   
#import mlflow
#mlflow.set_tracking_uri("http://localhost:5002")
#mlflow.set_experiment("train-trial")


"""
arguments
"""
args = fetch_args()
device = 'cuda' if torch.cuda.is_available() else 'cpu'
early_stopping = EarlyStopping(args['patience'], verbose=True, delta=args['delta'])


"""
loading data-set....
"""
print("==> loading data-set...")
train_loader, classes = gen_train_loader(args['train_path'], args['input_size'], args['train_batch_size'])
test_loader, _ = gen_test_loader(args['test_path'], args['input_size'], args['test_batch_size'])
num_classes = len(classes)
print(f'Total {num_classes} classes are: {classes}')


"""
model
"""
print("==> building model...")
net = VGG_19(num_classes=num_classes)
net = net.to(device)
summary(net, (3, 224, 224))


"""
training
"""
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=args['learning_rate'], momentum=0.9, weight_decay=5e-4)


def log_scalar(name, value, step):
    """Log a scalar value to both MLflow"""
    #mlflow.log_metric(name, value)


def train(epoch):
    print('\nEpoch:[%d/%d]' % (epoch, args['epochs']))
    net.train()
    loss, correct, total = 0, 0, 0

    for index, (inputs, targets) in enumerate(train_loader):
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = net(inputs)
        loss_curr = criterion(outputs, targets)
        loss_curr.backward()
        optimizer.step()

        loss += loss_curr.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()

        progress_bar(index, len(train_loader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
                     % (loss / (index + 1), 100. * correct / total, correct, total))

        log_scalar("train_loss", loss / (index + 1), epoch)

best_acc = 0

def test(epoch):
    global best_acc
    net.eval()
    loss, correct, total = 0, 0, 0

    with torch.no_grad():
        for index, (inputs, targets) in enumerate(test_loader):
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = net(inputs)
            loss_curr = criterion(outputs, targets)

            loss += loss_curr.item()
            _, predicted = outputs.max(1)
            print(targets)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()

            progress_bar(index, len(test_loader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
                         %(loss/(index+1), 100.*correct/total, correct, total))

    eval_loss = loss/(index+1)
    acc = 100. * correct / total
    if acc >= best_acc:
        print("Saving checkpoints...")
        state = {
            'net': net.state_dict(),
            'acc' : acc,
            'epoch' : epoch,
            'eval_loss' : eval_loss
        }
        if not os.path.isdir(args['ckpt_path']):
            os.mkdir(args['ckpt_path'])

        torch.save(state, args['ckpt_path'] + str('/ckpt_%d_acc%.2f.pt' % (epoch, acc)))
        best_acc = acc
    log_scalar("eval_loss", eval_loss, epoch)
    log_scalar("eval_acc", acc, epoch)

    return eval_loss


"""
with mlflow.start_run():
    # log parameters into mlflow
    mlflow.log_param("learning_rate", args['learning_rate'])
    mlflow.log_param("input_size", args['input_size'])
    mlflow.log_param("train_batch_size", args['train_batch_size'])
    mlflow.log_param("test_batch_size", args['test_batch_size'])
    mlflow.log_param("total_epochs", args['epochs'])
    mlflow.log_param("earlystop_patience", args['patience'])
    mlflow.log_param("earlystop_delta", args['delta'])

    for epoch in range(args['epochs']):
        train(epoch)
        eval_loss = test(epoch)

        # early stopping
        early_stopping(eval_loss, net)
        if early_stopping.early_stop:
            print("Early stopping.")
            break

    mlflow.log_artifacts(args['ckpt_path'])
"""



if __name__ == '__main__':
    for epoch in range(args['epochs']):
        train(epoch)
        eval_loss = test(epoch)

        # early stopping
        early_stopping(eval_loss, net)
        if early_stopping.early_stop:
            print("Early stopping.")
            break
# -*- coding:UTF-8 -*-
"""
@create: 2021-10-06 10:50:00
@author: GongChen
@description: 训练脚本
"""

import os
import torch
import torch.nn as nn
import torch.optim as optim
from torchsummary import summary

from models.vgg import *
from utils.arg_utils import *
from utils.data_utils import *
from utils.progress_utils import progress_bar
from utils.earlystopping import EarlyStopping   
#import mlflow
#mlflow.set_tracking_uri("http://localhost:5002")
#mlflow.set_experiment("train-trial")


"""
arguments
"""
args = fetch_args()
device = 'cuda' if torch.cuda.is_available() else 'cpu'
early_stopping = EarlyStopping(args['patience'], verbose=True, delta=args['delta'])


"""
loading data-set....
"""
print("==> loading data-set...")
train_loader, classes = gen_train_loader(args['train_path'], args['input_size'], args['train_batch_size'])
test_loader, _ = gen_test_loader(args['test_path'], args['input_size'], args['test_batch_size'])
num_classes = len(classes)
print(f'Total {num_classes} classes are: {classes}')


"""
model
"""
print("==> building model...")
net = VGG_19(num_classes=num_classes)
net = net.to(device)
summary(net, (3, 224, 224))


"""
training
"""
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=args['learning_rate'], momentum=0.9, weight_decay=5e-4)


def log_scalar(name, value, step):
    """Log a scalar value to both MLflow"""
    #mlflow.log_metric(name, value)


def train(epoch):
    print('\nEpoch:[%d/%d]' % (epoch, args['epochs']))
    net.train()
    loss, correct, total = 0, 0, 0

    for index, (inputs, targets) in enumerate(train_loader):
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = net(inputs)
        loss_curr = criterion(outputs, targets)
        loss_curr.backward()
        optimizer.step()

        loss += loss_curr.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()

        progress_bar(index, len(train_loader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
                     % (loss / (index + 1), 100. * correct / total, correct, total))

        log_scalar("train_loss", loss / (index + 1), epoch)

best_acc = 0

def test(epoch):
    global best_acc
    net.eval()
    loss, correct, total = 0, 0, 0

    with torch.no_grad():
        for index, (inputs, targets) in enumerate(test_loader):
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = net(inputs)
            loss_curr = criterion(outputs, targets)

            loss += loss_curr.item()
            _, predicted = outputs.max(1)
            print(targets)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()

            progress_bar(index, len(test_loader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
                         %(loss/(index+1), 100.*correct/total, correct, total))

    eval_loss = loss/(index+1)
    acc = 100. * correct / total
    if acc >= best_acc:
        print("Saving checkpoints...")
        state = {
            'net': net.state_dict(),
            'acc' : acc,
            'epoch' : epoch,
            'eval_loss' : eval_loss
        }
        if not os.path.isdir(args['ckpt_path']):
            os.mkdir(args['ckpt_path'])

        torch.save(state, args['ckpt_path'] + str('/ckpt_%d_acc%.2f.pt' % (epoch, acc)))
        best_acc = acc
    log_scalar("eval_loss", eval_loss, epoch)
    log_scalar("eval_acc", acc, epoch)

    return eval_loss


"""
with mlflow.start_run():
    # log parameters into mlflow
    mlflow.log_param("learning_rate", args['learning_rate'])
    mlflow.log_param("input_size", args['input_size'])
    mlflow.log_param("train_batch_size", args['train_batch_size'])
    mlflow.log_param("test_batch_size", args['test_batch_size'])
    mlflow.log_param("total_epochs", args['epochs'])
    mlflow.log_param("earlystop_patience", args['patience'])
    mlflow.log_param("earlystop_delta", args['delta'])

    for epoch in range(args['epochs']):
        train(epoch)
        eval_loss = test(epoch)

        # early stopping
        early_stopping(eval_loss, net)
        if early_stopping.early_stop:
            print("Early stopping.")
            break

    mlflow.log_artifacts(args['ckpt_path'])
"""



if __name__ == '__main__':
    for epoch in range(args['epochs']):
        train(epoch)
        eval_loss = test(epoch)

        # early stopping
        early_stopping(eval_loss, net)
        if early_stopping.early_stop:
            print("Early stopping.")
            break
# -*- coding:UTF-8 -*- @create: 2021-10-06 10:50:00 @author: GongChen @description: 训练脚本 #import mlflow #mlflow.set_tracking_uri("http://localhost:5002") #mlflow.set_experiment("train-trial") arguments loading data-set.... model training Log a scalar value to both MLflow #mlflow.log_metric(name, value) with mlflow.start_run():     # log parameters into mlflow     mlflow.log_param("learning_rate", args['learning_rate'])     mlflow.log_param("input_size", args['input_size'])     mlflow.log_param("train_batch_size", args['train_batch_size'])     mlflow.log_param("test_batch_size", args['test_batch_size'])     mlflow.log_param("total_epochs", args['epochs'])     mlflow.log_param("earlystop_patience", args['patience'])     mlflow.log_param("earlystop_delta", args['delta'])      for epoch in range(args['epochs']):         train(epoch)         eval_loss = test(epoch)          # early stopping         early_stopping(eval_loss, net)         if early_stopping.early_stop:             print("Early stopping.")             break      mlflow.log_artifacts(args['ckpt_path']) # early stopping