def train():
train_dataloader, val_dataloader = loadData()
pretrained_params = torch.load('VGG_pretrained.pth')
model = VGG()
# strict=False 使得预训练模型参数中和新模型对应上的参数会被载入，对应不上或没有的参数被抛弃。
model.load_state_dict(pretrained_params.state_dict(), strict=False)
if torch.cuda.is_available():
model.cuda()
# finetune 时冻结XXlayer的参数
# for p in model.XXlayers.parameters():
# p.requires_grad = False
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
loss_func = nn.CrossEntropyLoss()
best_acc = 0
for epoch in range(epochs):
epoch_loss = 0
steps = 0
for i, data in enumerate(train_dataloader):
inputs, labels = data
if torch.cuda.is_available():
inputs, labels = inputs.cuda(), labels.cuda()
inputs, labels = Variable(inputs), Variable(labels)
model.train()
optimizer.zero_grad()
outputs = model(inputs)
loss = loss_func(outputs, labels)
loss.backward()
optimizer.step()
epoch_loss += loss.data[0]
steps += 1
print('epoch:%d loss:%.3f' % (epoch + 1, epoch_loss / steps))
if epoch % 5 == 0:
val_acc = evaluate(model, val_dataloader)
if val_acc > best_acc:
best_acc = val_acc
torch.save(model, 'best_VGG.pkl')
torch.save(model.state_dict(), 'best_VGG_params.pkl')
print('test acc:'.format(val_acc))
print('Finished Training')
torch.save(model, 'VGG.pkl')
torch.save(model.state_dict(), 'VGG_params.pkl')
def train(train_loader, valid_loader, writer):
model = VGG(batch_size=cf.batch_size)
criterion = nn.BCELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=cf.learning_rate)
for e in range(cf.epoch):
batch_loss = 0
for batch_x, batch_y in train_loader:
y_predict = model(batch_x)
y_one_hot = to_convert_one_hot(y_predict)
loss = criterion(y_one_hot, batch_y.float())
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss = loss.item()
batch_loss += loss
print(f'batch loss: {loss:.3f}')
print(f'Epoch #{e}: --- Training loss: {batch_loss/cf.batch_size:.3f}')
writer.log_training(batch_loss / cf.batch_size, e)
save_path = './models/chkpt-%d.pt' % (e + 1)
torch.save(
{
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
'epoch': e
}, save_path)
print("Saved checkpoint to: %s" % save_path)
# iterate through all the features for the chosen layers
for gen_feature, orig_feature, style_feature in zip(
generated_features, original_img_features, style_features):
# batch_size will just be 1
batch_size, channel, height, width = gen_feature.shape
original_loss += torch.mean((gen_feature - orig_feature)**2)
# Compute Gram Matrix of generated
G = gen_feature.view(channel, height * width).mm(
gen_feature.view(channel, height * width).t())
# Compute Gram Matrix of Style
A = style_feature.view(channel, height * width).mm(
style_feature.view(channel, height * width).t())
style_loss += torch.mean((G - A)**2)
total_loss = alpha * original_loss + beta * style_loss
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
if step % 200 == 0:
print(total_loss)
temp = generated
temp = temp.to("cpu").squeeze(0)
to_pil_image = transforms.ToPILImage()
results = to_pil_image(temp)
#img = Image.show(temp)
results.save("generated" + str(step), "png")
torch.save(model.state_dict(),
"/home/abhrant/Neural_Style_Transfer_Web_App/nst_weight.pt")
num_batches += 1
if (count % 100 == 0):
# compute stats for the full training set
total_loss = running_loss / num_batches
total_error = running_error / num_batches
elapsed = (time.time() - start) / 60
if int(elapsed) % 10 == 0:
lost.append(total_loss)
errors.append(total_error)
torch.save(
{
'model': net.state_dict(),
'epoch': epoch,
'step': count,
'lr': my_lr,
'lost': lost,
'error': error
}, "../../models/base/vgg_base_model.pth")
print('epoch=', epoch, '\t step=', count, '/', step, '\t time=',
elapsed, 'min', '\t lr=', my_lr, '\t loss=', total_loss,
'\t error=', total_error * 100, 'percent')
#eval_on_test_set()
print(' ')
my_lr = my_lr / 2
def train(config):
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
device = torch.device(config.device)
random.seed(config.seed)
np.random.seed(config.seed)
torch.manual_seed(config.seed)
if config.use_gpu:
torch.cuda.manual_seed_all(config.seed)
# Dataset
infection_train = Infection_Dataset('train')
infection_test = Infection_Dataset('test')
infection_val = Infection_Dataset('val')
covid_train = Covid_Dataset('train')
covid_test = Covid_Dataset('test')
covid_val = Covid_Dataset('val')
# Dataloader from dataset
infection_train_loader = DataLoader(infection_train,
batch_size=config.batch_size,
shuffle=True)
infection_test_loader = DataLoader(infection_test,
batch_size=config.batch_size,
shuffle=True)
infection_val_loader = DataLoader(infection_val,
batch_size=config.batch_size,
shuffle=True)
covid_train_loader = DataLoader(covid_train,
batch_size=config.batch_size,
shuffle=True)
covid_test_loader = DataLoader(covid_test,
batch_size=config.batch_size,
shuffle=True)
covid_val_loader = DataLoader(covid_val,
batch_size=config.batch_size,
shuffle=True)
# L2 regularization parameter
l2_lambda = 0.001
# Instantiate model, criterion and oprimizer
model = VGG()
if config.use_gpu:
model.to(device)
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.AdamW(model.parameters(), lr=config.lr)
lr_sheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.7)
best_acc = -0.1
best_epoch = -1
start_time = time.time()
# Train the model that classifies normal and infection images
print("***** Training the first classifier *****")
for epoch in range(config.epochs):
total_loss = 0.0
model.train()
for images_data, target_labels in tqdm(infection_train_loader):
# images_data: [batch_size, 1, 150, 150]
# target_labels: [batch_size, 2]
if config.use_gpu:
images_data = images_data.cuda()
target_labels = target_labels.cuda()
total_loss = 0.0
model.train()
optimizer.zero_grad()
predicted_labels = model(images_data)
loss = criterion(predicted_labels, target_labels)
# L2 regularization
l2_reg = torch.tensor(0.)
if config.use_gpu:
l2_reg = l2_reg.cuda()
for param in model.parameters():
l2_reg += torch.norm(param)
loss += l2_lambda * l2_reg
loss.backward()
optimizer.step()
# Evaluate the performance and save the model parameters each epoch
train_acc, train_loss = evaluate(infection_train_loader, model)
val_acc, val_loss = evaluate(infection_test_loader, model)
torch.save(model.state_dict(),
'./checkpoints/' + str(epoch) + '_params_infection.pth')
# Save the best performing model parameters based on validation accuracy
if val_acc > best_acc:
best_acc = val_acc
best_epoch = epoch
torch.save(model.state_dict(),
'./checkpoints/' + 'best_params_infection.pth')
print(
f"{now()} Epoch{epoch}: train_loss: {train_loss}, val_loss: {val_loss}, train_acc: {train_acc}, val_acc: {val_acc}"
)
lr_sheduler.step()
# Record loss and accuracies for learning curve plots
fieldnames = [
'epoch', 'train_loss', 'val_loss', 'train_acc', 'val_acc'
]
out_dict = {
'epoch': epoch,
'train_loss': train_loss,
'val_loss': val_loss,
'train_acc': train_acc,
'val_acc': val_acc
}
with open('./outputs/infection.csv', 'a') as out_f:
writer = DictWriter(out_f, fieldnames=fieldnames)
writer.writerow(out_dict)
end_time = time.time()
print("*" * 20)
print(
f"{now()} finished; epoch {best_epoch} best_acc: {best_acc}, time/epoch: {(end_time-start_time)/config.epochs}"
)
print()
# Instantiate another model
model_covid = VGG()
if config.use_gpu:
model_covid.to(device)
optimizer_covid = optim.AdamW(model_covid.parameters(), lr=config.lr)
lr_sheduler_covid = optim.lr_scheduler.StepLR(optimizer_covid,
step_size=3,
gamma=0.7)
best_acc_covid = -0.1
best_epoch_covid = -1
start_time_covid = time.time()
# Train another model that classifies covid and non-covid images
print("***** Training the second classifier *****")
for epoch_covid in range(config.epochs):
total_loss_covid = 0.0
model_covid.train()
for images_data, target_labels in tqdm(covid_train_loader):
# images_data: [batch_size, 1, 150, 150]
# target_labels: [batch_size, 2]
if config.use_gpu:
images_data = images_data.cuda()
target_labels = target_labels.cuda()
total_loss_covid = 0.0
model_covid.train()
optimizer_covid.zero_grad()
predicted_labels = model_covid(images_data)
loss = criterion(predicted_labels, target_labels)
# L2 regularization
l2_reg = torch.tensor(0.)
if config.use_gpu:
l2_reg = l2_reg.cuda()
for param in model_covid.parameters():
l2_reg += torch.norm(param)
loss += l2_lambda * l2_reg
loss.backward()
optimizer_covid.step()
total_loss_covid += loss.item()
# Evaluate the performance and save the model parameters each epoch
train_acc_covid, train_loss_covid = evaluate(covid_train_loader,
model_covid)
val_acc_covid, val_loss_covid = evaluate(covid_test_loader,
model_covid)
torch.save(model_covid.state_dict(),
'./checkpoints/' + str(epoch_covid) + '_params_covid.pth')
# Save the best performing model parameters based on validation accuracy
if val_acc_covid > best_acc_covid:
best_acc_covid = val_acc_covid
best_epoch_covid = epoch_covid
torch.save(model_covid.state_dict(),
'./checkpoints/' + 'best_params_covid.pth')
print(
f"{now()} epoch {epoch_covid}: train_loss: {train_loss_covid}, val_loss: {val_loss_covid}, train_acc: {train_acc_covid}, val_acc_covid: {val_acc_covid}"
)
lr_sheduler_covid.step()
# Record loss and accuracies for learning curve plots
fieldnames = [
'epoch', 'train_loss', 'val_loss', 'train_acc', 'val_acc'
]
out_dict = {
'epoch': epoch_covid,
'train_loss': train_loss_covid,
'val_loss': val_loss_covid,
'train_acc': train_acc_covid,
'val_acc': val_acc_covid
}
with open('./outputs/infection.csv', 'a') as out_f:
writer = DictWriter(out_f, fieldnames=fieldnames)
writer.writerow(out_dict)
end_time = time.time()
print("*" * 20)
print(
f"{now()} finished; epoch {best_epoch_covid} best_acc_covid: {best_acc_covid}, time/epoch: {(end_time-start_time)/config.epochs}"
)
target = target.cuda()
output = net(images)
loss = loss_func(output, target)
opt.zero_grad()
loss.backward()
opt.step()
losses.append(loss.item())
print("Epoch: %4d, step: %4d, loss: %.3f" % (epoch, i, loss.item()))
mean_loss = np.mean(losses)
lr_decay.step()
writer.add_scalar('train/loss', mean_loss, epoch)
writer.add_scalar('learning rate', opt.param_groups[0]['lr'], epoch)
print("=============== Mean loss in epoch %4d is: %.3f ===============" % (epoch, mean_loss))
if epoch % 5 == 0:
precious, recall, accuracy = valid_accuracy(valid_loader, net)
print("=============== Epoch %4d - precious: %.3f, recall: %.3f, accuracy: %.3f ==============="
% (epoch, precious, recall, accuracy))
writer.add_scalars('validation', {'precious': precious,
'recall': recall,
'accuracy': accuracy}, epoch)
torch.save({'epoch': epoch, 'lr': opt.param_groups[0]['lr'], 'state_dict': net.state_dict()},
"./modules/vgg-%d-%.3f.pth.tar" % (epoch, mean_loss))
writer.close()
