train_dataloader = DataLoader(train_dataset, batch_size=8, drop_last=True)
val_dataloader = DataLoader(val_dataset, batch_size=8)
model = SiameseNetwork(backbone=args.backbone)
model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)
criterion = torch.nn.BCELoss()
writer = SummaryWriter(os.path.join(args.out_path, "summary"))
best_val = 10000000000
for epoch in range(args.epochs):
print("[{} / {}]".format(epoch, args.epochs))
model.train()
losses = []
correct = 0
total = 0
# Training Loop Start
for (img1, img2), y, (class1, class2) in train_dataloader:
img1, img2, y = map(lambda x: x.to(device), [img1, img2, y])
prob = model(img1, img2)
loss = criterion(prob, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
import argparse
from autolab_core import YamlConfig
from siamese import SiameseNetwork
if __name__ == "__main__":
# parse the provided configuration file, set tf settings, and train
conf_parser = argparse.ArgumentParser(description="Train Siamese model")
conf_parser.add_argument("--config", action="store", default="cfg/train.yaml",
dest="conf_file", type=str, help="path to the configuration file")
conf_args = conf_parser.parse_args()
# read in config file information from proper section
config = YamlConfig(conf_args.conf_file)
model_config = config['model']
train_config = config['train']
model = SiameseNetwork('training', model_config)
model.train(train_config)
device = torch.device('cuda:{}'.format(cuda_id))
net = net.to(device)
criterion = ContrastiveLoss(margin=1.0).cuda(device)
cudnn.benchmark = True
print('computation device: {}'.format(device))
def save_checkpoint(state, filename):
file_path = os.path.join(filename)
torch.save(state, file_path)
pdist = nn.PairwiseDistance(p=2)
for epoch in range(num_epoch):
net.train()
train_loader._sample_once()
running_loss = 0.0
running_num = 0
start = time.time()
positive_dist = 0.0
negative_dist = 0.0
for i in range(len(train_loader)):
x1, x2, label = train_loader[i]
x1, x2, label = x1.to(device), x2.to(device), label.to(device)
feat1, feat2 = net(x1, x2)