with open("Output.txt", "w") as text_file:
print(output_str, file=text_file)
if (epoch_conf_loss + epoch_loc_loss) < min_total_loss:
min_total_loss = epoch_conf_loss + epoch_loc_loss
best_model_wts = copy.deepcopy(model.state_dict())
torch.save(
ssd_net,
os.path.join(root_dir,
'ssd_net_1022_temporary_best_zoomout.pth'))
print("saved for temporary best model: ", idx_epochs)
print('Min Val Loss: {:4f}'.format(min_total_loss))
model.load_state_dict(best_model_wts)
with open("Output.txt", "w") as text_file:
print(output_str, file=text_file)
return model
num_classes = 3
root_dir = 'trained_model'
ssd_net = SSD(num_classes).cuda()
optimizer = optim.Adam(ssd_net.parameters())
criterion = MultiboxLoss()
ssd_net = train(ssd_net, optimizer, criterion, num_epochs=100)
torch.save(
ssd_net,
os.path.join(root_dir, 'ssd_net_1023_augmentation_100_zoomout_rgb.pth'))
batch_size=32,
shuffle=True,
num_workers=0)
# print('train items:', len(train_dataset))
idx, (bbox, label, img) = next(enumerate(train_data_loader))
# valid_dataset = csd.CityScapeDataset(train_list, False, False)
# valid_data_loader = torch.utils.data.DataLoader(valid_dataset,
# batch_size=4,
# shuffle=False,
# num_workers=0)
# print('validation items:', len(valid_dataset))
net = SSD(3)
optimizer = optim.SGD(net.parameters(),
lr=learning_rate,
momentum=0.9,
weight_decay=1e-4)
optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)
criterion = MultiboxLoss([0.1, 0.1, 0.2, 0.2])
if use_gpu:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
net.cuda()
criterion.cuda()
train_losses = []
valid_losses = []
itr = 0
for epoch_idx in range(0, max_epochs):
if WILL_TRAIN:
# Print key info of tensor inputs
idx, (loc_targets, conf_targets,
image) = next(enumerate(train_data_loader))
print('loc_targets tensor shape:', loc_targets.shape)
print('conf_targets tensor shape:', conf_targets.shape)
print('image tensor shape:', image.shape)
# Create the instance of our defined network
net = SSD(num_classes=3)
net.cuda()
criterion = MultiboxLoss([0.1, 0.1, 0.2, 0.2],
iou_threshold=0.5,
neg_pos_ratio=3.0)
# optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9, weight_decay=0.0005)
optimizer = torch.optim.Adam(net.parameters(), lr=0.001)
train_losses = []
valid_losses = []
max_epochs = 50
itr = 0
# Train process
for epoch_idx in range(0, max_epochs):
for train_batch_idx, (train_loc_targets, train_conf_targets,
train_image) in enumerate(train_data_loader):
itr += 1
net.train()
train_loc_targets = Variable(train_loc_targets.cuda().float())
train_conf_targets = Variable(train_conf_targets.cuda().long())
train_image = train_image.permute(0, 3, 1, 2)
train_image = Variable(train_image.cuda().float())