def main():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("using {} device.".format(device))
data_transform = {
"train":
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
]),
"val":
transforms.Compose([
transforms.Resize((224, 224)), # cannot 224, must (224, 224)
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
}
data_root = os.path.abspath(os.path.join(os.getcwd(),
"../")) # get data root path,脚本位置
image_path = os.path.join(data_root, "data_set") # rock data set path
assert os.path.exists(image_path), "{} path does not exist.".format(
image_path)
train_dataset = datasets.ImageFolder(root=os.path.join(
image_path, "train"),
transform=data_transform["train"])
train_num = len(train_dataset)
rock_list = train_dataset.class_to_idx
cla_dict = dict((val, key) for key, val in rock_list.items())
# write dict into json file
json_str = json.dumps(cla_dict, indent=4)
with open('class_indices.json', 'w') as json_file:
json_file.write(json_str)
batch_size = 32
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0,
8]) # number of workers
print('Using {} dataloader workers every process'.format(nw))
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=nw)
validate_dataset = datasets.ImageFolder(root=os.path.join(
image_path, "val"),
transform=data_transform["val"])
val_num = len(validate_dataset)
validate_loader = torch.utils.data.DataLoader(validate_dataset,
batch_size=4,
shuffle=False,
num_workers=nw)
print("using {} images for training, {} images for validation.".format(
train_num, val_num))
net = AlexNet(num_classes=7, init_weights=True)
net.to(device)
loss_function = nn.CrossEntropyLoss()
# pata = list(net.parameters())
optimizer = optim.Adam(net.parameters(), lr=0.0002)
epochs = 10
save_path = './AlexNet.pth'
best_acc = 0.0
train_steps = len(train_loader)
for epoch in range(epochs):
# train
net.train()
running_loss = 0.0
train_bar = tqdm(train_loader)
for step, data in enumerate(train_bar):
images, labels = data
optimizer.zero_grad()
outputs = net(images.to(device))
loss = loss_function(outputs, labels.to(device))
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(
epoch + 1, epochs, loss)
# validate
net.eval()
acc = 0.0 # accumulate accurate number / epoch
with torch.no_grad():
val_bar = tqdm(validate_loader)
for val_data in val_bar:
val_images, val_labels = val_data
outputs = net(val_images.to(device))
predict_y = torch.max(outputs, dim=1)[1]
acc += torch.eq(predict_y, val_labels.to(device)).sum().item()
val_accurate = acc / val_num
print('[epoch %d] train_loss: %.3f val_accuracy: %.3f' %
(epoch + 1, running_loss / train_steps, val_accurate))
if val_accurate > best_acc:
best_acc = val_accurate
torch.save(net.state_dict(), save_path)
print('Finished Training')
# 可视化
writer = SummaryWriter(
comment='_alexnet_go_Adam_lr={}_momentum={}_epochs={}'.format(
lr, momentum, epochs))
# 模型
model = AlexNet(num_classes=100)
model.to('cuda:0')
writer.add_graph(model, torch.randn([256, 3, 224, 224]).cuda())
# 损失
loss_func = torch.nn.CrossEntropyLoss()
# 优化器 SGD(loss func, grad func...)
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad,
model.parameters()),
weight_decay=weight_decay,
momentum=momentum,
lr=lr)
# optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=lr, betas=(0.9, 0.99))
# 学习率更新
lr_schedule = lambda epoch: ((epoch < epochs * 0.8) * lr * 0.95 +
(epoch >= epochs * 0.8) * lr * 0.8)
# 数据集
train_loader = dataset.get_aug_dataloader(image_folder, batch_size=128)
for epoch in range(epochs):
# model.train() ：启用BatchNormalization和Dropout, model.eval() ：不启用BatchNormalization和Dropout