def validate(model, val_path, img_path, mask_path, batch_size):
running_metrics_val = Score(NUM_CLASSES)
val_loss_meter = AverageMeter()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.eval()
val_dataset = PascalVOCDataset(list_file=val_path,
img_dir=img_path,
mask_dir=mask_path,
img_size=512,
is_transform=True)
val_dataloader = DataLoader(val_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=4)
class_weights = 1.0 / val_dataset.get_class_probability().to(device)
criterion = torch.nn.CrossEntropyLoss(weight=class_weights).to(device)
with torch.no_grad():
for batch_idx, (image, mask) in tqdm.tqdm(enumerate(val_dataloader), desc='Evaluating...'):
input_tensor = torch.autograd.Variable(image.to(device))
target_tensor = torch.autograd.Variable(mask.to(device))
output = model(input_tensor)
loss = criterion(output, target_tensor)
pred = output.data.max(1)[1].cpu().numpy()
gt = target_tensor.data.cpu().numpy()
running_metrics_val.update(gt, pred)
val_loss_meter.update(loss.item())
score, class_iou = running_metrics_val.get_scores()
return val_loss_meter.avg, score, class_iou
# Set absolute target data path
mask_dir = os.path.join(data_root, args.mask_dir)
# Set model path
SAVED_MODEL_PATH = args.model_path
# Set output directory path
OUTPUT_DIR = args.output_dir
# Set whether to use cuda
CUDA = args.gpu is not None
# Set GPU ID
GPU_ID = args.gpu
# Initialize Pascal VOC Dataset for validation
val_dataset = PascalVOCDataset(list_file=val_path,
img_dir=img_dir,
mask_dir=mask_dir)
# Initialize validation dataloader
val_dataloader = DataLoader(
val_dataset,
batch_size=BATCH_SIZE, # mini-batch size
shuffle=True, # Use shuffle
num_workers=4) # number of cpu
# If using cuda
if CUDA:
# Initialize SegNet model on gpu memory
model = SegNet(input_channels=NUM_INPUT_CHANNELS,
output_channels=NUM_OUTPUT_CHANNELS).cuda(GPU_ID)
print("Epoch #{}\tLoss: {:.8f}\t Time: {:2f}s".format(epoch+1, loss_f, delta))
if __name__ == "__main__":
data_root = args.data_root
train_path = os.path.join(data_root, args.train_path)
img_dir = os.path.join(data_root, args.img_dir)
mask_dir = os.path.join(data_root, args.mask_dir)
CUDA = args.gpu is not None
GPU_ID = args.gpu
train_dataset = PascalVOCDataset(list_file=train_path,
img_dir=img_dir,
mask_dir=mask_dir)
train_dataloader = DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=4)
if args.model == "unet":
model = UNet(input_channels=NUM_INPUT_CHANNELS,
output_channels=NUM_OUTPUT_CHANNELS)
elif args.model == "segnet":
model = SegNet(input_channels=NUM_INPUT_CHANNELS,
output_hannels=NUM_OUTPUT_CHANNELS)
else:
os.makedirs('checkpoints', exist_ok=True)
data_root = args.data_root
train_path = os.path.join(data_root, args.train_path)
train_img_path = os.path.join(data_root, args.img_dir)
train_mask_path = os.path.join(data_root, args.mask_dir)
val_path = os.path.join(data_root, args.val_path)
val_img_path = os.path.join(data_root, args.img_dir)
val_mask_path = os.path.join(data_root, args.mask_dir)
writer = Writer('logs')
logger = log('')
train_dataset = PascalVOCDataset(list_file=train_path,
img_dir=train_img_path,
mask_dir=train_mask_path,
img_size=512,
is_transform=True)
train_dataloader = DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=4)
model = SegNet().to(device)
class_weights = 1.0 / train_dataset.get_class_probability()
print(class_weights)
criterion = torch.nn.CrossEntropyLoss(weight=class_weights).to(device)
# start from checkpoint
if args.checkpoint:
def train():
# NOTE: Change the file paths here appropriately
if args.container:
homedir = '/scratch/scratch1/sidnayak/'
else:
homedir = os.path.expanduser("~")
data_root = homedir + "/data/VOCdevkit/VOC2007"
list_file_path = os.path.join(data_root, "ImageSets", "Main", "train.txt")
img_dir = os.path.join(data_root, "JPEGImages")
if not os.path.exists('./Weights'):
os.makedirs('./Weights')
dataset = PascalVOCDataset(list_file_path, img_dir, args.orig_shape,
args.target_shape)
trainloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True)
if args.rm_runs:
shutil.rmtree('runs')
if args.tensorboard:
print('Init tensorboardX')
writer = SummaryWriter(log_dir='runs/{}'.format(
datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")))
# CUDA compatability
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
print('#' * 50)
print(device)
AE1 = AutoEncoder(args.orig_shape, args.latent_shape).to(device)
AE2 = AutoEncoder(args.target_shape, args.latent_shape).to(device)
recon_criterion = getCriterion(args.recon_loss_type)
latent_criterion = getCriterion(args.latent_loss_type)
optimizer_1 = torch.optim.Adam(AE1.parameters(),
lr=args.lr,
weight_decay=1e-5)
optimizer_2 = torch.optim.Adam(AE2.parameters(),
lr=args.lr,
weight_decay=1e-5)
i = 0
for epoch in range(args.num_epochs):
for batch in trainloader:
orig_env_image = torch.autograd.Variable(
batch['orig_env_image']).to(device)
target_env_image = torch.autograd.Variable(
batch['target_env_image']).to(device)
z1, s1 = AE1(orig_env_image)
z2, s2 = AE2(target_env_image)
if args.scale_loss:
reconstruction_loss1 = recon_criterion(
orig_env_image, s1) / (args.orig_shape**2)
reconstruction_loss2 = recon_criterion(
target_env_image, s2) / (args.target_shape**2)
latent_loss = latent_criterion(z1, z2) / args.latent_shape
loss = reconstruction_loss1 + reconstruction_loss2 + latent_loss
else:
reconstruction_loss1 = recon_criterion(orig_env_image, s1)
reconstruction_loss2 = recon_criterion(target_env_image, s2)
latent_loss = latent_criterion(z1, z2)
loss = reconstruction_loss1 + reconstruction_loss2 + latent_loss
if args.tensorboard:
writer.add_scalar('Autoencoder_1_Loss',
reconstruction_loss1.item(), i)
writer.add_scalar('Autoencoder_2_Loss',
reconstruction_loss2.item(), i)
writer.add_scalar('Latent_Loss', latent_loss.item(), i)
writer.add_scalar('Total_Loss', loss.item(), i)
optimizer_1.zero_grad()
loss.backward(retain_graph=True)
optimizer_1.step()
optimizer_2.zero_grad()
loss.backward()
optimizer_2.step()
i += 1
print(loss.item(), reconstruction_loss1.item(),
reconstruction_loss2.item(), latent_loss.item())
print('Saving Weights')
torch.save(
{
'model_state_dict': AE1.state_dict(),
'optimizer_state_dict': optimizer_1.state_dict(),
}, args.weight_paths[0])
torch.save(
{
'model_state_dict': AE2.state_dict(),
'optimizer_state_dict': optimizer_2.state_dict(),
}, args.weight_paths[1])
def test():
# NOTE: Change the file paths here appropriately
if args.container:
homedir = '/scratch/scratch1/sidnayak/'
else:
homedir = os.path.expanduser("~")
data_root = homedir + "/data/VOCdevkit/VOC2007"
list_file_path = os.path.join(data_root, "ImageSets", "Main", "train.txt")
img_dir = os.path.join(data_root, "JPEGImages")
if not os.path.exists('./Weights'):
os.makedirs('./Weights')
dataset = PascalVOCDataset(list_file_path, img_dir, args.orig_shape,
args.target_shape)
trainloader = DataLoader(dataset, batch_size=1, shuffle=True)
AE1 = AutoEncoder(args.orig_shape, args.latent_shape)
AE2 = AutoEncoder(args.target_shape, args.latent_shape)
criterion = nn.MSELoss()
if os.path.exists(args.weight_paths[0]):
print('Loading weights')
checkpoint1 = torch.load(args.weight_paths[0], map_location='cpu')
checkpoint2 = torch.load(args.weight_paths[1], map_location='cpu')
AE1.load_state_dict(checkpoint1['model_state_dict'])
AE2.load_state_dict(checkpoint2['model_state_dict'])
AE1.eval()
AE2.eval()
i = 0
for batch in trainloader:
orig_env_image = torch.autograd.Variable(batch['orig_env_image'])
target_env_image = torch.autograd.Variable(batch['target_env_image'])
z1, s1 = AE1(orig_env_image)
z2, s2 = AE2(target_env_image)
img1 = tensor2img(orig_env_image)
img2 = tensor2img(target_env_image)
s1_img = tensor2img(s1.detach())
s2_img = tensor2img(s2.detach())
plt.figure()
ax = plt.subplot(2, 2, 1)
im = ax.imshow(img1)
plt.title('64x64')
ax = plt.subplot(2, 2, 2)
im = ax.imshow(s1_img)
plt.title('reconstruct(64x64)')
ax = plt.subplot(2, 2, 3)
im = ax.imshow(img1)
plt.title('32x32')
ax = plt.subplot(2, 2, 4)
im = ax.imshow(s2_img)
plt.title('reconstruct(32x32)')
plt.tight_layout()
print(list(z1.detach().numpy()))
print(list(z2.detach().numpy()))
# print(criterion(z1,z2))
plt.show()
# plot z1 and z2
plt.plot(np.array(list(z1.detach().numpy()))[0], label='z1')
plt.plot(np.array(list(z2.detach().numpy()))[0], label='z2')
plt.grid()
plt.legend()
plt.show()
i += 1
if i == 1:
break
import tensorflow as tf
from tensorflow import keras
from dataset import PascalVOCDataset, get_data_list
from model import SSD300
model = SSD300()
train_imgs, train_anns, val_imgs, val_anns, test_imgs, tes_anns = get_data_list(
)
train_dataset = PascalVOCDataset(split='train',
images=train_imgs,
annotations=train_anns,
batch_size=1)
opt = keras.optimizers.Adam(1e-4)
for i in range(2000):
imgs, true_locs, true_cls = next(train_dataset)
with tf.GradientTape() as tape:
pred_locs, pred_cls = model(imgs)
loss, _, _, _ = model.loss(true_locs, true_cls, pred_locs, pred_cls)
gradients = tape.gradient(loss, model.trainable_variables)
opt.apply_gradients(zip(gradients, model.trainable_variables))
if i % 20 == 0:
print(i, ' LOSS:', loss.numpy())
model.save('ssd300.h5')
img_dir = os.path.join(data_root, args.img_dir)
mask_dir = os.path.join(data_root, args.mask_dir)
CUDA = args.gpu is not None
print("CUDA", CUDA)
GPU_ID = args.gpu
print('GPU', GPU_ID)
print('DATASET')
if args.custom:
train_dataset = ScoreDataset(list_file=train_path,
img_dir=img_dir,
mask_dir=mask_dir)
else:
train_dataset = PascalVOCDataset(list_file=train_path,
img_dir=img_dir,
mask_dir=mask_dir)
print('LOADER')
train_dataloader = DataLoader(train_dataset,
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=4)
if CUDA:
model = SegNet(input_channels=NUM_INPUT_CHANNELS,
output_channels=NUM_OUTPUT_CHANNELS).cuda(GPU_ID)
class_weights = 1.0 / train_dataset.get_class_probability().cuda(
GPU_ID)
criterion = torch.nn.CrossEntropyLoss(
weight=class_weights).cuda(GPU_ID)