-
Notifications
You must be signed in to change notification settings - Fork 0
/
experiment_train.py
111 lines (88 loc) · 5.28 KB
/
experiment_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
from src.model import *
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
def train_model(model, num_epochs, dataloaders, lr=args.lr, momentum=args.momentum, beta2=args.beta2, log_interval=100, **kwargs):
model.to(device)
if beta2 > 0.:
optimizer = torch.optim.Adam(model.parameters(), lr=lr, betas=(momentum, beta2)) #, amsgrad=amsgrad)
else:
optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum) # to set training variables
df_train = pd.DataFrame([], columns=['epoch', 'avg_loss', 'avg_acc', 'avg_loss_val', 'avg_acc_val', 'device_type'])
for epoch in range(num_epochs):
loss_train = 0
acc_train = 0
for i, (images, labels) in enumerate(dataloaders['train']):
images, labels = images.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs[:,0], labels.float())
loss.backward()
optimizer.step()
loss_train += loss.item() * images.size(0)
preds = torch.round(torch.sigmoid(outputs[:,0].data))
acc_train += torch.sum(preds == labels.data)
avg_loss = loss_train / dataset_sizes[task]['train']
avg_acc = acc_train / dataset_sizes[task]['train']
with torch.no_grad():
loss_val = 0
acc_val = 0
for i, (images, labels) in enumerate(dataloaders['val']):
images, labels = images.to(device), labels.to(device)
outputs = model(images)
loss = criterion(outputs[:,0], labels.float())
loss_val += loss.item() * images.size(0)
preds = torch.round(torch.sigmoid(outputs[:,0].data))
acc_val += torch.sum(preds == labels.data)
avg_loss_val = loss_val / dataset_sizes[task]['val']
avg_acc_val = acc_val / dataset_sizes[task]['val']
df_train.loc[epoch] = {'epoch':epoch, 'avg_loss':avg_loss, 'avg_acc':float(avg_acc),
'avg_loss_val':avg_loss_val, 'avg_acc_val':float(avg_acc_val), 'device_type':device.type}
print(f"Epoch {epoch+1}/{num_epochs} : train= loss: {avg_loss:.4f} / acc : {avg_acc:.4f} - val= loss : {avg_loss_val:.4f} / acc : {avg_acc_val:.4f}")
model.cpu()
torch.cuda.empty_cache()
return model, df_train
# https://towardsdatascience.com/cross-entropy-for-classification-d98e7f974451
criterion = nn.BCEWithLogitsLoss() #binary_cross_entropy_with_logits
# Training and saving the network
models_vgg = {}
opt = {}
n_output = 1
# Downloading the model
model_filenames = {}
for model_name in all_models:
task = 'animal' if 'animal' in model_name else 'artifact'
model_filenames[model_name] = args.model_path + model_name + '.pt'
filename = f'results/{datetag}_{args.HOST}_train_{model_name}.json'
models_vgg[model_name] = torchvision.models.vgg16(pretrained=True)
num_features = models_vgg[model_name].classifier[-1].in_features
features = list(models_vgg[model_name].classifier.children())[:-1] # Remove last layer
features.extend([nn.Linear(num_features, n_output)]) # Add our layer with 10 outputs
models_vgg[model_name].classifier = nn.Sequential(*features) # Replace the model classifier
if os.path.isfile(model_filenames[model_name]):
print("Loading pretrained model for..", model_name, ' from', model_filenames[model_name])
if device.type == 'cuda':
models_vgg[model_name].load_state_dict(torch.load(model_filenames[model_name])) #on GPU
else:
models_vgg[model_name].load_state_dict(torch.load(model_filenames[model_name], map_location=torch.device('cpu'))) #on CPU
else:
print("Re-training pretrained model...", model_filenames[model_name])
since = time.time()
p = 1 if 'gray' in model_name else 0
if 'scale' in model_name:
df_train = None
for image_size_ in args.image_sizes: # starting with low resolution images
print(f"Traning {model_name}, image_size = {image_size_}, p (Grayscale) = {p}")
(dataset_sizes, dataloaders, image_datasets, data_transforms) = datasets_transforms(image_size=image_size_, p=p)
models_vgg[model_name], df_train_ = train_model(models_vgg[model_name], num_epochs=args.num_epochs//len(args.image_sizes),
dataloaders=dataloaders[task])
df_train = df_train_ if df_train is None else df_train.append(df_train_, ignore_index=True)
else :
print(f"Traning {model_name}, image_size = {args.image_size}, p (Grayscale) = {p}")
(dataset_sizes, dataloaders, image_datasets, data_transforms) = datasets_transforms(image_size=args.image_size, p=p)
models_vgg[model_name], df_train = train_model(models_vgg[model_name], num_epochs=args.num_epochs,
dataloaders=dataloaders[task])
torch.save(models_vgg[model_name].state_dict(), model_filenames[model_name])
df_train.to_json(filename)
elapsed_time = time.time() - since
print(f"Training completed in {elapsed_time // 60:.0f}m {elapsed_time % 60:.0f}s")
print()