import torchvision
import torchvision.transforms as transforms
from datasets import camvid
import utils.imgs
import utils.training as train_utils
from datasets import joint_transforms
from pathlib import Path
from models import tiramisu
CAMVID_PATH = Path('/home/jingwenlai/data', 'CamVid/CamVid')
batch_size = 2
normalize = transforms.Normalize(mean=camvid.mean, std=camvid.std)
test_dset = camvid.CamVid(CAMVID_PATH,
'test',
joint_transform=None,
transform=transforms.Compose(
[transforms.ToTensor(), normalize]))
test_loader = torch.utils.data.DataLoader(test_dset,
batch_size=batch_size,
shuffle=False)
print("Test: %d" % len(test_loader.dataset.imgs))
model = tiramisu.FCDenseNet67(n_classes=12).cuda()
model_weights = ".weights/latest.th"
startEpoch = train_utils.load_weights(model, model_weights)
print("load_weights, return epoch: ", startEpoch)
train_utils.view_sample_predictions(model, test_loader, n=10)
val_dataset = Dataset(os.path.join(args.valdir, 'Image/'),
os.path.join(args.valdir, 'LabeledImage/'), mytransforms)
val_loader = DataLoader(val_dataset, batch_size=1, shuffle=True, num_workers=1)
model = FCDenseNet103(args.classes)
#model=model.cuda()
model = nn.DataParallel(model).cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-5)
#optimizer = torch.optim.RMSprop(model.parameters(), lr=args.lr, weight_decay=1e-4)
criterion = nn.NLLLoss2d()
train_utils.load_weights(model, 'weights/weights-0-0.178-0.000.pth')
for epoch in range(args.epochs):
print('\n _______________________________________________')
trn_loss, trn_err = train_utils.train(model, train_loader, optimizer,
criterion, epoch)
print('Epoch {:d}\nTrain - Loss: {:.4f}'.format(epoch, trn_loss))
train_utils.save_weights(model, epoch, float(trn_loss), 0)
### Adjust Lr ###
train_utils.adjust_learning_rate(args.lr, args.decay, optimizer, epoch,
DECAY_EVERY_N_EPOCHS)
### Validate ###
train_utils.view_sample_predictions(model, val_loader)
def test():
train_utils.test(model, test_loader, criterion, epoch=1)
train_utils.view_sample_predictions(model, test_loader, n=1)