-
Notifications
You must be signed in to change notification settings - Fork 1
/
train.py
85 lines (76 loc) · 2.63 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
from model import DogModel
from data import DogDataset
from tqdm import tqdm
import torch
from torch import nn, optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import transforms
def distance(x1, x2):
return 1.0 - F.cosine_similarity(x1, x2, dim=1)
def check_grads(params):
grad = 0
i = 0
for p in params:
if p.grad is not None:
grad += p.grad.norm()
i += 1
return grad / i
def train():
epochs = 20
batch_size = 32
train_transform = transforms.Compose([
# transforms.RandomAffine(15, translate=(0.2, 0.2), scale=(0.5, 1.5), shear=15),
transforms.RandomResizedCrop((128, 128)),
transforms.ToTensor(),
])
test_transform = transforms.Compose([
transforms.Resize((128, 128)),
transforms.ToTensor(),
])
train_dataset = DogDataset(transform=train_transform, train=True)
test_dataset = DogDataset(transform=test_transform, train=False)
trainloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
testloader = DataLoader(test_dataset, batch_size=batch_size)
device = "cuda" if torch.cuda.is_available() else "cpu"
net = DogModel(metric_size=2)
net.to(device)
criterion = nn.TripletMarginWithDistanceLoss(distance_function=distance, margin=0.5)
optimizer = optim.Adam(net.parameters(), lr=0.0001, weight_decay=1e-5)
for ep in range(epochs):
train_loss = 0
net.train()
for batch in tqdm(trainloader):
x, pos, neg, label = batch
x = x.to(device)
pos = pos.to(device)
neg = neg.to(device)
x = net(x)
pos = net(pos)
neg = net(neg)
loss = criterion(x, pos, neg)
optimizer.zero_grad()
loss.backward()
optimizer.step()
train_loss += loss.item()
print(f"Epoch {ep*1}: Train loss: {train_loss / len(trainloader)}")
net.eval()
test_loss = 0
for batch in tqdm(testloader):
x, pos, neg, label = batch
with torch.no_grad():
x = x.to(device)
pos = pos.to(device)
neg = neg.to(device)
x = net(x)
pos = net(pos)
neg = net(neg)
loss = criterion(x, pos, neg)
test_loss += loss.item()
print(f"Epoch {ep*1}: Test loss: {test_loss / len(testloader)}")
torch.save({
"net": net.state_dict(),
"optimzer": optimizer.state_dict()
}, "net.pt")
if __name__ == "__main__":
train()