rename trial.py to hello.py.
This commit is contained in:
Brainway
Gitee
279
hello.py
Normal file
279
hello.py
Normal file
@ -0,0 +1,279 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from vit_pytorch import ViT,SimpleViT
|
||||
|
||||
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 = 600,
|
||||
patch_size = 30,
|
||||
num_classes = 5,
|
||||
dim = 768,
|
||||
depth = 12,
|
||||
heads = 12,
|
||||
mlp_dim = 3072,
|
||||
dropout = 0.1,
|
||||
emb_dropout = 0.1
|
||||
)
|
||||
|
||||
# model = SimpleViT(
|
||||
# image_size = 600,
|
||||
# patch_size = 30,
|
||||
# num_classes = 2,
|
||||
# dim = 256,
|
||||
# depth = 6,
|
||||
# heads = 16,
|
||||
# mlp_dim = 256
|
||||
# )
|
||||
|
||||
# model = ViT(
|
||||
# #Vit-best
|
||||
# # image_size = 600,
|
||||
# # patch_size = 30,
|
||||
# # num_classes = 5,
|
||||
# # dim = 512,
|
||||
# # depth = 6,
|
||||
# # heads = 8,
|
||||
# # mlp_dim = 512,
|
||||
# # pool = 'cls',
|
||||
# # channels = 3,
|
||||
# # dim_head = 12,
|
||||
# # dropout = 0.1,
|
||||
# # emb_dropout = 0.1
|
||||
|
||||
# #Vit-small
|
||||
# image_size = 600,
|
||||
# patch_size = 30,
|
||||
# num_classes = 5,
|
||||
# dim = 256,
|
||||
# depth = 8,
|
||||
# heads = 16,
|
||||
# mlp_dim = 256,
|
||||
# pool = 'cls',
|
||||
# channels = 3,
|
||||
# dim_head = 16,
|
||||
# dropout = 0.1,
|
||||
# emb_dropout = 0.1
|
||||
|
||||
# #Vit-tiny
|
||||
# # image_size = 600,
|
||||
# # patch_size = 30,
|
||||
# # num_classes = 5,
|
||||
# # dim = 256,
|
||||
# # depth = 4,
|
||||
# # heads = 6,
|
||||
# # mlp_dim = 256,
|
||||
# # pool = 'cls',
|
||||
# # channels = 3,
|
||||
# # dim_head = 6,
|
||||
# # 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)
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user