228 lines
6.5 KiB
Python
228 lines
6.5 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
from vit_pytorch import ViT
|
||
|
||
from utils.data_utils import get_loader_new
|
||
from utils.scheduler import WarmupCosineSchedule
|
||
|
||
from tqdm import tqdm
|
||
|
||
import os
|
||
import numpy as np
|
||
|
||
def net():
|
||
model = ViT(
|
||
image_size = 320,
|
||
patch_size = 32,
|
||
num_classes = 5,
|
||
dim = 768,
|
||
depth = 4,
|
||
heads = 12,
|
||
mlp_dim = 1024,
|
||
pool = 'cls',
|
||
channels = 3,
|
||
dim_head = 12,
|
||
dropout = 0.1,
|
||
emb_dropout = 0.1
|
||
)
|
||
|
||
|
||
return model
|
||
|
||
|
||
|
||
#计算模型参数数量
|
||
def count_parameters(model):
|
||
params = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||
return params/1000000
|
||
|
||
|
||
#Loss平均
|
||
class AverageMeter(object):
|
||
"""Computes and stores the average and current value"""
|
||
def __init__(self):
|
||
self.reset()
|
||
|
||
def reset(self):
|
||
self.val = 0
|
||
self.avg = 0
|
||
self.sum = 0
|
||
self.count = 0
|
||
|
||
def update(self, val, n=1):
|
||
self.val = val
|
||
self.sum += val * n
|
||
self.count += n
|
||
self.avg = self.sum / self.count
|
||
|
||
criterion = nn.CrossEntropyLoss()
|
||
|
||
#简单准确率
|
||
def simple_accuracy(preds, labels):
|
||
return (preds == labels).mean()
|
||
|
||
#模型测试
|
||
def test(device, model, test_loader, global_step):
|
||
eval_losses = AverageMeter()
|
||
|
||
print("***** Running Validation *****")
|
||
|
||
|
||
model.eval()
|
||
all_preds, all_label = [], []
|
||
epoch_iterator = tqdm(test_loader,
|
||
desc="Validating... (loss=X.X)",
|
||
bar_format="{l_bar}{r_bar}",
|
||
dynamic_ncols=True)
|
||
|
||
for step, batch in enumerate(epoch_iterator):
|
||
batch = tuple(t.to(device) for t in batch)
|
||
x, y = batch
|
||
with torch.no_grad():
|
||
logits = model(x)
|
||
|
||
eval_loss = criterion(logits, y)
|
||
eval_loss = eval_loss.mean()
|
||
eval_losses.update(eval_loss.item())
|
||
|
||
preds = torch.argmax(logits, dim=-1)
|
||
|
||
if len(all_preds) == 0:
|
||
all_preds.append(preds.detach().cpu().numpy())
|
||
all_label.append(y.detach().cpu().numpy())
|
||
else:
|
||
all_preds[0] = np.append(
|
||
all_preds[0], preds.detach().cpu().numpy(), axis=0
|
||
)
|
||
all_label[0] = np.append(
|
||
all_label[0], y.detach().cpu().numpy(), axis=0
|
||
)
|
||
epoch_iterator.set_description("Validating... (loss=%2.5f)" % eval_losses.val)
|
||
|
||
all_preds, all_label = all_preds[0], all_label[0]
|
||
accuracy = simple_accuracy(all_preds, all_label)
|
||
accuracy = torch.tensor(accuracy).to(device)
|
||
|
||
val_accuracy = accuracy.detach().cpu().numpy()
|
||
|
||
|
||
print("test Loss: %2.5f" % eval_losses.avg)
|
||
print("test Accuracy: %2.5f" % val_accuracy)
|
||
|
||
return val_accuracy
|
||
|
||
|
||
#保存模型
|
||
def save_model(model):
|
||
model_checkpoint = os.path.join('./output', "%s_vit_checkpoint.pth" % 'ieemooempty')
|
||
torch.save(model, model_checkpoint)
|
||
print("Saved model checkpoint to [DIR: %s]", './output')
|
||
|
||
|
||
#训练
|
||
def train(model,train_loader,device,train_NUM_STEPS,LEARNING_RATE,WEIGHT_DECAY,WARMUP_STEPS,test_loader):
|
||
|
||
optimizer = torch.optim.SGD(model.parameters(),
|
||
lr=LEARNING_RATE,
|
||
momentum=0.9,
|
||
weight_decay=WEIGHT_DECAY)
|
||
t_total = train_NUM_STEPS
|
||
scheduler = WarmupCosineSchedule(optimizer, warmup_steps=WARMUP_STEPS, t_total=t_total)
|
||
model.zero_grad()
|
||
losses = AverageMeter()
|
||
global_step, best_acc = 0, 0
|
||
gradient_accumulation_steps = 1
|
||
while True:
|
||
model.train()
|
||
epoch_iterator = tqdm(train_loader,
|
||
desc="Training (X / X Steps) (loss=X.X)",
|
||
bar_format="{l_bar}{r_bar}",
|
||
dynamic_ncols=True)
|
||
all_preds, all_label = [], []
|
||
for step, batch in enumerate(epoch_iterator):
|
||
batch = tuple(t.to(device) for t in batch)
|
||
x, y = batch
|
||
logits = model(x)
|
||
|
||
loss = criterion(logits,y)
|
||
|
||
loss.backward()
|
||
|
||
preds = torch.argmax(logits, dim=-1)
|
||
|
||
if len(all_preds) == 0:
|
||
all_preds.append(preds.detach().cpu().numpy())
|
||
all_label.append(y.detach().cpu().numpy())
|
||
else:
|
||
all_preds[0] = np.append(
|
||
all_preds[0], preds.detach().cpu().numpy(), axis=0
|
||
)
|
||
all_label[0] = np.append(
|
||
all_label[0], y.detach().cpu().numpy(), axis=0
|
||
)
|
||
|
||
if (step + 1) % gradient_accumulation_steps == 0:
|
||
losses.update(loss.item()*gradient_accumulation_steps)
|
||
|
||
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
|
||
scheduler.step()
|
||
optimizer.step()
|
||
optimizer.zero_grad()
|
||
global_step += 1
|
||
|
||
epoch_iterator.set_description(
|
||
"Training (%d / %d Steps) (loss=%2.5f)" % (global_step, t_total, losses.val)
|
||
)
|
||
|
||
model.train() #需要2次，才会保存训练好的模型
|
||
|
||
|
||
if global_step % t_total == 0:
|
||
|
||
|
||
break
|
||
|
||
all_preds, all_label = all_preds[0], all_label[0]
|
||
accuracy = simple_accuracy(all_preds, all_label)
|
||
accuracy = torch.tensor(accuracy).to(device)
|
||
train_accuracy = accuracy.detach().cpu().numpy()
|
||
print("train accuracy: %f" % train_accuracy)
|
||
accuracy = test(device, model, test_loader, global_step)
|
||
|
||
if best_acc < accuracy:
|
||
save_model(model)
|
||
best_acc = accuracy
|
||
|
||
losses.reset()
|
||
if global_step % t_total == 0:
|
||
break
|
||
|
||
|
||
|
||
|
||
if __name__ == "__main__":
|
||
|
||
model = net()
|
||
train_loader, test_loader = get_loader_new()
|
||
trainnumsteps = len(train_loader)
|
||
testnumsteps = len(test_loader)
|
||
|
||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
|
||
print(str(count_parameters(model))+'MB')
|
||
# img = torch.randn(1, 3, 320, 320) #(batchsize,channels,width,higth) #3072是默认channels为3，3*1024
|
||
# preds = model(img)
|
||
# print(preds)
|
||
|
||
model = model.to(device)
|
||
epoch = 300
|
||
train_NUM_STEPS = trainnumsteps * epoch
|
||
LEARNING_RATE = 3e-2
|
||
WEIGHT_DECAY = 0
|
||
WARMUP_STEPS = 500
|
||
train(model,train_loader,device,train_NUM_STEPS,LEARNING_RATE,WEIGHT_DECAY,WARMUP_STEPS,test_loader)
|
||
|
||
|
||
|
||
|