classNotMinstDataset(Dataset): def__init__(self, file_dir): self.imgs = [] labels = [] # 从所有的文件夹中读取图片 for root, sub_folders, files in os.walk(file_dir): for name in files: self.imgs.append(os.path.join(root, name)) labels.append(root.split("\\")[1])
import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torch.utils.data import Dataset, DataLoader import torchvision.transforms as transforms from PIL import Image from sklearn import preprocessing import os import time
defforward(self, x): x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2)) x = F.max_pool2d(F.relu(self.conv2(x)), (2, 2)) x = x.view(x.size(0), -1) x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x
net = Net().to(device) criterion = nn.CrossEntropyLoss().to(device) optimizer = optim.Adam(net.parameters(), lr=0.001)
trainloader = data(True) testloader = data(False) t0 = time.time() # training for epoch inrange(1): running_loss = 0.0 for batch_idx, (input, label) inenumerate(trainloader): input, label = input.to(device), label.to(device).squeeze() optimizer.zero_grad() output = net(input) loss = criterion(output, label) loss.backward() optimizer.step()