-
Notifications
You must be signed in to change notification settings - Fork 0
/
vgg_cifar10_image_classification.py
executable file
·121 lines (90 loc) · 3.65 KB
/
vgg_cifar10_image_classification.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
from vgg import VGG
import torch
import torch.optim as optim
import torch.nn as nn
from torchvision import datasets, transforms
from utils import save_model, load_yaml
# Set the configuration
config = load_yaml("./config/vgg_cifar10_config.yml")
# Training setting
device = 'cuda' if torch.cuda.is_available() else 'cpu'
torch.manual_seed(config['data']['seed'])
if device == 'cuda':
torch.cuda.manual_seed_all(config['data']['seed'])
# Set the transform
transform = transforms.Compose([transforms.ToTensor(),
transforms.Resize(config['data']['img_size'])])
# Set the training data
train_data = datasets.CIFAR10(config['data']['data_path'], download=config['data']['download'], train=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_data, batch_size=config['data']['batch_size'], shuffle=config['data']['shuffle'], drop_last=config['data']['drop_last'])
# Set the test data
test_data = datasets.CIFAR10(config['data']['data_path'], download=config['data']['download'], train=False, transform=transform)
test_loader = torch.utils.data.DataLoader(train_data, batch_size=config['data']['batch_size'], shuffle=config['data']['shuffle'], drop_last=config['data']['drop_last'])
# Set the model
model = VGG(input_img=config['model']['input_img'], output_channel=config['model']['output_channel'], num_blocks=config['model']['num_blocks']).to(device)
print(model, device)
# Set the criterion and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(),
lr=config['train']['lr'],
betas=config['train']['betas'],
eps=config['train']['eps'],
weight_decay=config['train']['weight_decay'])
# Training
def train(epoch, data_loader, optimizer, criterion):
model.train()
train_loss = 0.0
train_num = 0
for i, data in enumerate(data_loader):
# get the inputs; data is a list of [inputs, labels]
inputs, labels = data
# Transfer data to device
inputs = inputs.to(device)
labels = labels.to(device)
# Model inference
outputs = model(inputs)
# Training
optimizer.zero_grad()
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# loss
train_loss += loss.item()
train_num += labels.size(0)
if i % config['others']['log_period'] == 0 and i != 0:
print(f'[{epoch}, {i}]\t Train loss: {train_loss / train_num:.10f}')
# Average loss
train_loss /= train_num
return train_loss
# Validation
def test(data_loader):
model.eval()
corrects = 0
test_num = 0
for data in data_loader:
# get the inputs; data is a list of [inputs, labels]
inputs, labels = data
# Transfer data to device
inputs = inputs.to(device)
labels = labels.to(device)
# Model inference
outputs = model(inputs)
# Calculate the accuracy
preds = torch.argmax(outputs.data, 1)
corrects += torch.sum(preds == labels.data).item()
# Number of the data
test_num += labels.size(0)
# Test accuracy
test_accuracy = 100 * corrects / test_num
return test_accuracy
# Main
if __name__ == '__main__':
for epoch in range(config['train']['epochs']): # loop over the dataset multiple times
# Training
train_loss = train(epoch, train_loader, optimizer, criterion)
# Validation
test_accuracy = test(test_loader)
# Print the log
print(f'Epoch: {epoch}\t Train loss: {train_loss:.10f}\t Test accuracy: {test_accuracy:.3f}')
# Save the model
save_model(model_name=config['save']['model_name'], epoch=epoch, model=model, optimizer=optimizer, loss=train_loss, config=config)