import torch import torchvision.models as models from torch.utils import data from torch import nn from torch import optim import numpy as np import argparse from data.MyDataset import Mydatasetpro from data.MyDataset import all_imgs_path, all_labels, transform

ROOT = "D:\UC_data\Softmax" DATA_PATH = ROOT + "\data\img"

parser = argparse.ArgumentParser(description='Soft') parser.add_argument('--name', type=str, default='resnet50', metavar='N', help='model name') parser.add_argument('--data_dir', type=str, default=DATA_PATH, help='path for datasets') parser.add_argument('--batch_size', type=int, default=32, metavar='N', help='input batch size for training') parser.add_argument('--epochs', type=int, default=1000, metavar='N', help='number of epochs to train') parser.add_argument('--lr', type=float, default=1e-3, metavar='LR', help='learning rate')

args = parser.parse_args() print(args)

==================================================================

Set Device

==================================================================

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

==================================================================

Set Model

==================================================================

model = models.resnet50(pretrained=True)

==================================================================

Divide the data

==================================================================

my_dataset = Mydatasetpro(all_imgs_path, all_labels, transform) my_datalodaer = data.DataLoader( my_dataset, batch_size=args.batch_size, shuffle=True)

imgs_batch, labels_batch = next(iter(my_datalodaer)) print(imgs_batch.shape)

划分数据集和测试集

index = np.random.permutation(len(all_imgs_path))

all_imgs_path = np.array(all_imgs_path)[index] all_labels = np.array(all_labels)[index]

80%做训练集

s = int(len(all_imgs_path) * 0.8)

train_imgs = all_imgs_path[:s] train_labels = all_labels[:s] test_imgs = all_imgs_path[s:] test_labels = all_labels[s:]

train_ds = Mydatasetpro(train_imgs, train_labels, transform) # TrainSet TensorData test_ds = Mydatasetpro(test_imgs, test_labels, transform) # TestSet TensorData

print(train_ds)

print(test_ds)

train_dl = data.DataLoader(train_ds, batch_size=args.batch_size, shuffle=True) # TrainSet Labels test_dl = data.DataLoader(train_ds, batch_size=args.batch_size, shuffle=True) # TestSet Labels

in_features = model.fc.in_features model.fc = nn.Sequential( nn.Linear(in_features, 256), nn.ReLU(), # nn.Dropout(0, 4), nn.Linear(256,3), nn.LogSoftmax(dim=1))

将模型迁移到gpu

model = model.to(DEVICE)

优化器

loss_fn = nn.CrossEntropyLoss() loss_fn = loss_fn.to(DEVICE) # 将loss_fn迁移到GPU

Adam损失函数

optimizer = optim.Adam(model.fc.parameters(), lr=args.lr)

steps = 0 running_loss = 0 count = 10 train_losses, test_losses = [], [] best_accuracy = 0.0 # 保存最佳模型的准确率

for epoch in range(args.epochs): model.train() # 遍历训练集数据 for imgs, labels in my_datalodaer: steps += 1 labels = torch.tensor(labels, dtype=torch.long) imgs, labels = imgs.to(DEVICE), labels.to(DEVICE) optimizer.zero_grad() # 梯度归零 outputs = model(imgs) loss = loss_fn(outputs, labels) loss.backward() # 反向传播计算梯度 optimizer.step() # 梯度优化 running_loss += loss.item()

    if steps % count == 0:
        test_loss = 0
        accuracy = 0
        model.eval()
        with torch.no_grad():
            # 遍历测试集数据
            for imgs, labels in test_dl:
                labels = torch.tensor(labels, dtype=torch.long)
                imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)
                outputs = model(imgs)
                loss = loss_fn(outputs, labels)
                test_loss += loss.item()
                ps = torch.exp(outputs)
                top_p, top_class = ps.topk(1, dim=1)
                equals = top_class == labels.view(*top_class.shape)
                accuracy += torch.mean(equals.type(torch.FloatTensor)).item()
        train_losses.append(running_loss / len(train_dl))
        test_losses.append(test_loss / len(test_dl))
        if steps % (int(50)) == 0:
            print(
                f"Epoch {epoch + 1}/{args.epochs}.. "
                f"Train loss: {running_loss / count:.3f}.. "
                f"Test loss: {test_loss / len(test_dl):.3f}.. "
                f"Test accuracy: {accuracy / len(test_dl):.3f}")
        running_loss = 0
        model.train()
        if accuracy > best_accuracy:
            best_accuracy = accuracy
            torch.save(model, "Direction_model.pth")
        if len(test_losses) > 50 and len(set(test_losses[-50:])) == 1:
            # 当最佳准确率有50个epoch不改变时候更新学习率
            for param_group in optimizer.param_groups:
                param_group['lr'] *= 0.1

print("Training completed."