基于pytorch的训练、保存模型、调用模型的四种天气分类

import torch
import torch.nn as nn
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import torch.nn.functional as F
import torchvision
import time
import os
import shutil
import copy
from torchvision import datasets,transforms
from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader
from sklearn.model_selection import train_test_split

start=time.time()

base_dir = r'D:/pycharmworkspace/ISLR-master/fourweather'

if not os.path.isdir(base_dir):
    os.mkdir(base_dir)
    train_dir = os.path.join(base_dir,'train')
    test_dir = os.path.join(base_dir, 'test')
    os.mkdir(train_dir)
    os.mkdir(test_dir)

specises=['cloudy','rain','shine','sunrise']

for train_or_test in ['train','test']:
        for spc in specises:
            exist = os.path.join(base_dir,train_or_test,spc)
            if not os.path.isdir(exist):
                os.mkdir(os.path.join(base_dir,train_or_test,spc))

image_dir = r'D:/pycharmworkspace/ISLR-master/4w'
for i,img in enumerate(os.listdir(image_dir)):
    for spec in specises:
        if spec in img:
            s = os.path.join(image_dir,img)
            if i%5 == 0:
                d = os.path.join(base_dir,'test',spec,img)
            else:
                d = os.path.join(base_dir, 'train',spec,img)
            shutil.copy(s,d)

transformation = transforms.Compose([
    transforms.Resize((96,96)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5,0.5,0.5],std=[0.5,0.5,0.5])
])

train_ds = torchvision.datasets.ImageFolder(
    os.path.join(base_dir,'train'),
    transform=transformation
)
test_ds = torchvision.datasets.ImageFolder(
    os.path.join(base_dir, 'test'),
    transform=transformation
)

BATCHSIZE=64

train_dl = torch.utils.data.DataLoader(
    train_ds,
    batch_size=BATCHSIZE,
    shuffle=True,
)
test_dl = torch.utils.data.DataLoader(
    train_ds,
    batch_size=BATCHSIZE,
)

imgs,labels=next(iter(train_dl))

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3,16,3,)
        self.bn1 = nn.BatchNorm2d(16)
        self.conv2 = nn.Conv2d(16,32,3,)
        self.bn2 = nn.BatchNorm2d(32)
        self.conv3 = nn.Conv2d(32,64,3,)
        self.bn3 = nn.B