Main

In [ ]:

import argparse
from tools import load,train,show
from models import LeNet,AlexNet,ResNet,VGGNet,MobileNet,GoogLeNet
from torchvision import transforms
import torch

import argparse from tools import load,train,show from models import LeNet,AlexNet,ResNet,VGGNet,MobileNet,GoogLeNet from torchvision import transforms import torch

In [ ]:

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, default='LeNet')
    parser.add_argument('--epoch', default=7, type=int)
    parser.add_argument('--batch_size', default=128, type=int)
    parser.add_argument('--lr', default=1e-4, type=float)
    parser.add_argument('--dropout', default=0.5, type=float)
    parser.add_argument('--device',default='cuda',type=str)
    args = parser.parse_args()
    if torch.cuda.is_available()==False and args.device=='cuda':
        args.device='cpu'

    data_path = "data"
    if args.model == 'VGGNet':
        args.batch_size=64 #否则容易出现cuda out of memory
    
    train_data, valid_data, test_data = load.load(data_path,batch_size=args.batch_size,val_rate=0.2)

    if args.model == 'LeNet':
        model = LeNet.LeNet()
        pic_shape = (32, 32)
    elif args.model == 'AlexNet':
        model = AlexNet.AlexNet(dropout_p=args.dropout)
        pic_shape = (227, 227)
    elif args.model == 'ResNet':
        model = ResNet.ResNet()
        pic_shape = (224, 224)
    elif args.model == 'VGGNet':
        model = VGGNet.VGGNet(dropout=args.dropout)
        pic_shape = (64, 64)
    elif args.model == 'MobileNet':
        model = MobileNet.MobileNet()
        pic_shape = (128, 128)
    elif args.model == 'GoogLeNet':
        model = GoogLeNet.GoogLeNet()
        pic_shape = (128, 128)


    '''
    loss_histories,acc_histories = [],[]
    for lr in [1e-4,1e-3,0.01,0.1]:
        optimizer = torch.optim.Adam(model.parameters(), lr=lr)
        torch_resize = transforms.Resize(pic_shape)
        loss_history,acc_history=train.train(model, optimizer, train_data, valid_data, args.device, args.epoch, torch_resize)
        loss_history = [tensor.item() for tensor in loss_history]
        acc_history = [tensor.item() for tensor in acc_history]
        loss_histories.append(loss_history)
        acc_histories.append(acc_history)

    show.draw2(loss_histories,acc_histories,[1e-4,1e-3,0.01,0.1])

    '''
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
    torch_resize = transforms.Resize(pic_shape)
    loss_history,acc_history=train.train(model, optimizer, train_data, valid_data, args.device, args.epoch, torch_resize)

    loss_history = [tensor.item() for tensor in loss_history]
    acc_history = [tensor.item() for tensor in acc_history]

    #print(loss_history,acc_history)
    #show.draw1(loss_history,acc_history)

    loss_sum, test_acc = train.get_acc(model, test_data, args.device, torch_resize)
    print(f"Test loss {loss_sum} acc {test_acc}%")

if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--model', type=str, default='LeNet') parser.add_argument('--epoch', default=7, type=int) parser.add_argument('--batch_size', default=128, type=int) parser.add_argument('--lr', default=1e-4, type=float) parser.add_argument('--dropout', default=0.5, type=float) parser.add_argument('--device',default='cuda',type=str) args = parser.parse_args() if torch.cuda.is_available()==False and args.device=='cuda': args.device='cpu' data_path = "data" if args.model == 'VGGNet': args.batch_size=64 #否则容易出现cuda out of memory train_data, valid_data, test_data = load.load(data_path,batch_size=args.batch_size,val_rate=0.2) if args.model == 'LeNet': model = LeNet.LeNet() pic_shape = (32, 32) elif args.model == 'AlexNet': model = AlexNet.AlexNet(dropout_p=args.dropout) pic_shape = (227, 227) elif args.model == 'ResNet': model = ResNet.ResNet() pic_shape = (224, 224) elif args.model == 'VGGNet': model = VGGNet.VGGNet(dropout=args.dropout) pic_shape = (64, 64) elif args.model == 'MobileNet': model = MobileNet.MobileNet() pic_shape = (128, 128) elif args.model == 'GoogLeNet': model = GoogLeNet.GoogLeNet() pic_shape = (128, 128) ''' loss_histories,acc_histories = [],[] for lr in [1e-4,1e-3,0.01,0.1]: optimizer = torch.optim.Adam(model.parameters(), lr=lr) torch_resize = transforms.Resize(pic_shape) loss_history,acc_history=train.train(model, optimizer, train_data, valid_data, args.device, args.epoch, torch_resize) loss_history = [tensor.item() for tensor in loss_history] acc_history = [tensor.item() for tensor in acc_history] loss_histories.append(loss_history) acc_histories.append(acc_history) show.draw2(loss_histories,acc_histories,[1e-4,1e-3,0.01,0.1]) ''' optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) torch_resize = transforms.Resize(pic_shape) loss_history,acc_history=train.train(model, optimizer, train_data, valid_data, args.device, args.epoch, torch_resize) loss_history = [tensor.item() for tensor in loss_history] acc_history = [tensor.item() for tensor in acc_history] #print(loss_history,acc_history) #show.draw1(loss_history,acc_history) loss_sum, test_acc = train.get_acc(model, test_data, args.device, torch_resize) print(f"Test loss {loss_sum} acc {test_acc}%")