Hey Thomas, I tried what you suggested and the outcome looks like the given plot
class T_LEAP(nn.Module):
"""T_LEAP ARCHITECTURE"""
def __init__(self):
super(T_LEAP, self).__init__()
self.encoder = torch.nn.Sequential(
torch.nn.Conv3d(3, 64, kernel_size=(3,3,3), stride=(1,1,1), padding=(1,1,1)),
torch.nn.BatchNorm3d(num_features=64),
torch.nn.ReLU(inplace=True),
torch.nn.MaxPool3d(kernel_size=(1,2,2),stride=(1,2,2)),
torch.nn.Conv3d(64, 128, kernel_size=(3,3,3), stride=(1,1,1), padding=(1,1,1)),
torch.nn.BatchNorm3d(num_features=128),
torch.nn.ReLU(inplace=True),
torch.nn.MaxPool3d(kernel_size=(2,2,2),stride=(2,2,2)),
torch.nn.Conv3d(128, 256, kernel_size=(3,3,3), stride=(1,1,1), padding=(1,1,1)),
torch.nn.BatchNorm3d(num_features=256),
torch.nn.ReLU(inplace=True),
torch.nn.MaxPool3d(kernel_size=(1,2,2),stride=(1,2,2)),
torch.nn.Conv3d(256, 512, kernel_size=(1,3,3), stride=(1,1,1), padding=(1,1,1)),
torch.nn.BatchNorm3d(num_features=512),
torch.nn.ReLU(inplace=True),
)
self.decoder = torch.nn.Sequential(
torch.nn.ConvTranspose3d(512, 256, kernel_size=(1,3,3), stride = (1,2,2), padding=(1,1,1), output_padding=(0,1,1)),
torch.nn.BatchNorm3d(num_features=256),
torch.nn.ReLU(inplace=True),
torch.nn.Conv3d(256, 256, kernel_size=(1,3,3), stride=(1,1,1), padding=(1,1,1)),
torch.nn.BatchNorm3d(num_features=256),
torch.nn.ReLU(inplace=True),
torch.nn.ConvTranspose3d(256,128, kernel_size=(1,3,3), stride = (1,2,2), padding=(1,1,1),output_padding=(0,1,1)),
torch.nn.BatchNorm3d(num_features=128),
torch.nn.ReLU(inplace=True),
torch.nn.Conv3d(128, 128, kernel_size=(1,3,3), stride=(1,1,1), padding=(1,1,1)),
torch.nn.BatchNorm3d(num_features=128),
torch.nn.ReLU(inplace=True),
torch.nn.ConvTranspose3d(128,16, kernel_size=(1,3,3), stride = (1,2,2), padding=(1,1,1),output_padding=(0,1,1))
)
def forward(self, image):
"""PUTTING THE MODEL TOGETHER"""
encoder = self.encoder(image)
decoder = self.decoder(encoder)
decoder_shape = decoder.shape
flatten_decoder = decoder.view(decoder_shape[0], decoder_shape[1], -1)
flatten_heat_map = torch.nn.Softmax(dim=2)(flatten_decoder )
return flatten_heat_map.view(decoder_shape)
model = T_LEAP()
if torch.cuda.is_available():
input = torch.rand(1,3, 2, 256, 256).cuda()
model = model.cuda()
print("Output image shape:",model(input).shape)
else:
summary(model,
input_size=(3, 2, 200, 200),
batch_size=1
)
print("************************************************")
# torch.cuda.memory_summary(device=None, abbreviated=False)
Sample training loop:
#Get parameters, start training model
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
EPOCHS = 10
lr = 1e-3
model = T_LEAP()
criterion = torch.nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(),lr=lr, amsgrad=True)
model = model.to(DEVICE)
# def softXEnt(input, target):
# logprobs = torch.nn.functional.log_softmax (input, dim = 1)
# return -(target * logprobs).mean()
tr_losses =[]
val_losses=[]
def loss_plot(epochs, train, val):
ep = [i for i in range(epochs)]
plt.plot(ep,train,label="Training loss")
plt.plot(ep,val,label="Validation loss")
plt.title("Training and validation loss")
plt.show()
image = torch.rand(1,3,2,160, 160).to(DEVICE)
heatmap = torch.rand(1,16,1,160,160).to(DEVICE)
for idx in range(EPOCHS):
model.train()
# Get images and transfer to GPU
image = image
heatmap = heatmap
optimizer.zero_grad()
output = model(image)
loss = criterion(output,heatmap)
loss.backward()
optimizer.step()
curr_trloss = loss.item()
# Print losses
print(f"Training loss {idx}:",curr_trloss/1)
# Evaluation loop
model.eval()
with torch.no_grad():
# Get images and transfer to GPU
image = image
heatmap = heatmap
output = model(image)
loss = criterion(output,heatmap)
curr_valoss = loss.item()
print(f"Validation loss {idx}:",curr_valoss/1)
print(" ")
tr_losses.append(curr_trloss/1)
val_losses.append(curr_valoss/1)
print("Training and evalutation is now complete")
loss_plot(EPOCHS,tr_losses, val_losses)
Output
yellow - validation curve (sorry about the legends)
