-
Notifications
You must be signed in to change notification settings - Fork 180
/
train.py
148 lines (118 loc) · 5.56 KB
/
train.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
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
# camera-ready
import sys
from datasets import DatasetTrain, DatasetVal # (this needs to be imported before torch, because cv2 needs to be imported before torch for some reason)
sys.path.append("/root/deeplabv3/model")
from deeplabv3 import DeepLabV3
sys.path.append("/root/deeplabv3/utils")
from utils import add_weight_decay
import torch
import torch.utils.data
import torch.nn as nn
from torch.autograd import Variable
import torch.optim as optim
import torch.nn.functional as F
import numpy as np
import pickle
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import cv2
import time
# NOTE! NOTE! change this to not overwrite all log data when you train the model:
model_id = "1"
num_epochs = 1000
batch_size = 3
learning_rate = 0.0001
network = DeepLabV3(model_id, project_dir="/root/deeplabv3").cuda()
train_dataset = DatasetTrain(cityscapes_data_path="/root/deeplabv3/data/cityscapes",
cityscapes_meta_path="/root/deeplabv3/data/cityscapes/meta")
val_dataset = DatasetVal(cityscapes_data_path="/root/deeplabv3/data/cityscapes",
cityscapes_meta_path="/root/deeplabv3/data/cityscapes/meta")
num_train_batches = int(len(train_dataset)/batch_size)
num_val_batches = int(len(val_dataset)/batch_size)
print ("num_train_batches:", num_train_batches)
print ("num_val_batches:", num_val_batches)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=batch_size, shuffle=True,
num_workers=1)
val_loader = torch.utils.data.DataLoader(dataset=val_dataset,
batch_size=batch_size, shuffle=False,
num_workers=1)
params = add_weight_decay(network, l2_value=0.0001)
optimizer = torch.optim.Adam(params, lr=learning_rate)
with open("/root/deeplabv3/data/cityscapes/meta/class_weights.pkl", "rb") as file: # (needed for python3)
class_weights = np.array(pickle.load(file))
class_weights = torch.from_numpy(class_weights)
class_weights = Variable(class_weights.type(torch.FloatTensor)).cuda()
# loss function
loss_fn = nn.CrossEntropyLoss(weight=class_weights)
epoch_losses_train = []
epoch_losses_val = []
for epoch in range(num_epochs):
print ("###########################")
print ("######## NEW EPOCH ########")
print ("###########################")
print ("epoch: %d/%d" % (epoch+1, num_epochs))
############################################################################
# train:
############################################################################
network.train() # (set in training mode, this affects BatchNorm and dropout)
batch_losses = []
for step, (imgs, label_imgs) in enumerate(train_loader):
#current_time = time.time()
imgs = Variable(imgs).cuda() # (shape: (batch_size, 3, img_h, img_w))
label_imgs = Variable(label_imgs.type(torch.LongTensor)).cuda() # (shape: (batch_size, img_h, img_w))
outputs = network(imgs) # (shape: (batch_size, num_classes, img_h, img_w))
# compute the loss:
loss = loss_fn(outputs, label_imgs)
loss_value = loss.data.cpu().numpy()
batch_losses.append(loss_value)
# optimization step:
optimizer.zero_grad() # (reset gradients)
loss.backward() # (compute gradients)
optimizer.step() # (perform optimization step)
#print (time.time() - current_time)
epoch_loss = np.mean(batch_losses)
epoch_losses_train.append(epoch_loss)
with open("%s/epoch_losses_train.pkl" % network.model_dir, "wb") as file:
pickle.dump(epoch_losses_train, file)
print ("train loss: %g" % epoch_loss)
plt.figure(1)
plt.plot(epoch_losses_train, "k^")
plt.plot(epoch_losses_train, "k")
plt.ylabel("loss")
plt.xlabel("epoch")
plt.title("train loss per epoch")
plt.savefig("%s/epoch_losses_train.png" % network.model_dir)
plt.close(1)
print ("####")
############################################################################
# val:
############################################################################
network.eval() # (set in evaluation mode, this affects BatchNorm and dropout)
batch_losses = []
for step, (imgs, label_imgs, img_ids) in enumerate(val_loader):
with torch.no_grad(): # (corresponds to setting volatile=True in all variables, this is done during inference to reduce memory consumption)
imgs = Variable(imgs).cuda() # (shape: (batch_size, 3, img_h, img_w))
label_imgs = Variable(label_imgs.type(torch.LongTensor)).cuda() # (shape: (batch_size, img_h, img_w))
outputs = network(imgs) # (shape: (batch_size, num_classes, img_h, img_w))
# compute the loss:
loss = loss_fn(outputs, label_imgs)
loss_value = loss.data.cpu().numpy()
batch_losses.append(loss_value)
epoch_loss = np.mean(batch_losses)
epoch_losses_val.append(epoch_loss)
with open("%s/epoch_losses_val.pkl" % network.model_dir, "wb") as file:
pickle.dump(epoch_losses_val, file)
print ("val loss: %g" % epoch_loss)
plt.figure(1)
plt.plot(epoch_losses_val, "k^")
plt.plot(epoch_losses_val, "k")
plt.ylabel("loss")
plt.xlabel("epoch")
plt.title("val loss per epoch")
plt.savefig("%s/epoch_losses_val.png" % network.model_dir)
plt.close(1)
# save the model weights to disk:
checkpoint_path = network.checkpoints_dir + "/model_" + model_id +"_epoch_" + str(epoch+1) + ".pth"
torch.save(network.state_dict(), checkpoint_path)