Hello i need help with measuring Iou. Can someone help me or give example how can i save best model during the training of this model ?
class Model(torch.nn.Module):
def __init__(self, input_shape= (3, 448, 640), output_shape= (2, 448, 640)):
super(Model, self).__init__()
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.layers_encoder_0 = [
self.conv_bn(input_shape[0], 64, 2),
self.conv_bn(64, 64, 1),
self.conv_bn(64, 128, 2),
self.conv_bn(128, 128, 1),
self.conv_bn(128, 256, 2)
]
self.layers_encoder_1 = [
self.conv_bn(256, 256, 1),
self.conv_bn(256, 256, 2),
self.conv_bn(256, 256, 1),
self.conv_bn(256, 256, 2)
]
self.layers_decoder = [
self.conv_bn(256 + 256, 256, 1),
self.conv_bn(256, 128, 1),
self.conv_bn(128, 128, 1),
nn.Conv2d(128, output_shape[0], kernel_size=1, stride=1, padding=0),
nn.Upsample(scale_factor=8, mode="bilinear", align_corners=False)
]
for i in range(len(self.layers_encoder_0)):
if hasattr(self.layers_encoder_0[i], "weight"):
torch.nn.init.xavier_uniform_(self.layers_encoder_0[i].weight)
for i in range(len(self.layers_encoder_1)):
if hasattr(self.layers_encoder_1[i], "weight"):
torch.nn.init.xavier_uniform_(self.layers_encoder_1[i].weight)
for i in range(len(self.layers_decoder)):
if hasattr(self.layers_decoder[i], "weight"):
torch.nn.init.xavier_uniform_(self.layers_decoder[i].weight)
self.model_encoder_0 = nn.Sequential(*self.layers_encoder_0)
self.model_encoder_0.to(self.device)
self.model_encoder_1 = nn.Sequential(*self.layers_encoder_1)
self.model_encoder_1.to(self.device)
self.model_decoder = nn.Sequential(*self.layers_decoder)
self.model_decoder.to(self.device)
print(self.model_encoder_0)
print(self.model_encoder_1)
print(self.model_decoder)
def forward(self, x):
encoder_0 = self.model_encoder_0(x)
encoder_1 = self.model_encoder_1(encoder_0)
encoder_1_up = torch.nn.functional.interpolate(encoder_1, scale_factor=4, mode="nearest")
d_in = torch.cat([encoder_0, encoder_1_up], dim=1)
y = self.model_decoder(d_in)
return y
def save(self, path):
torch.save(self.model_encoder_0.state_dict(), path + "model_encoder_0.pt")
torch.save(self.model_encoder_1.state_dict(), path + "model_encoder_1.pt")
torch.save(self.model_decoder.state_dict(), path + "model_decoder.pt")
def load(self, path):
self.model_encoder_0.load_state_dict(torch.load(path + "model_encoder_0.pt", map_location=self.device))
self.model_encoder_1.load_state_dict(torch.load(path + "model_encoder_1.pt", map_location=self.device))
self.model_decoder.load_state_dict(torch.load(path + "model_decoder.pt", map_location=self.device))
self.model_encoder_0.eval()
self.model_encoder_1.eval()
self.model_decoder.eval()
def conv_bn(self, inputs, outputs, stride):
return nn.Sequential(
nn.Conv2d(inputs, outputs, kernel_size=3, stride=stride, padding=1),
nn.BatchNorm2d(outputs),
nn.ReLU(inplace=True))
def tconv_bn(self, inputs, outputs, stride):
return nn.Sequential(
nn.ConvTranspose2d(inputs, outputs, kernel_size=3, stride=stride, padding=1, output_padding=1),
nn.BatchNorm2d(outputs),
nn.ReLU(inplace=True))