Hello,
I’m training a VAE using a database of satellite images (256x256). Here is some of my code:
class VAE(nn.Module):
def __init__(self, inp_s, conv_kernel_size1):
super(VAE, self).__init__()
self.n1=1
self.n2=20
self.n3=50
self.padding=0
self.stride=1
self.size_out_conv1=int(((inp_s[2]-conv_kernel_size1-2*self.padding)/self.stride)+1)
self.size_out_conv2=int(((self.size_out_conv1-conv_kernel_size1-2*self.padding)/self.stride)+1)
self.fc1 = nn.Linear(self.size_out_conv2*self.size_out_conv2*self.n3, 700)
self.fc21 = nn.Linear(700, 20)
self.fc22 = nn.Linear(700, 20)
self.fc3 = nn.Linear(20, 700)
self.fc4 = nn.Linear(700, self.size_out_conv2*self.size_out_conv2*self.n3)
self.conv1=nn.Conv2d(self.n1, self.n2, kernel_size=conv_kernel_size1)
self.conv2=nn.Conv2d(self.n2, self.n3, kernel_size=conv_kernel_size1)
self.deconv1=nn.ConvTranspose2d(self.n2, self.n1, kernel_size=conv_kernel_size1)
self.deconv2=nn.ConvTranspose2d(self.n3, self.n2, kernel_size=conv_kernel_size1)
def encode(self, x):
x=F.relu(self.conv1(x))
x=F.relu(self.conv2(x))
s=x.shape
x=x.view(-1, s[1]*s[2]*s[3])
h1 =F.relu(self.fc1(x))
h1 =F.dropout(F.relu(self.fc1(x)), training=self.training)
return self.fc21(h1), self.fc22(h1), s
def reparameterize(self, mu, logvar):
std = torch.exp(0.5*logvar)
eps = torch.randn_like(std)
return mu + eps*std
def decode(self, z, s):
h3 = F.relu(self.fc3(z))
x=F.relu(self.fc4(h3))
x=x.view(s[0],s[1],s[2],s[3])
x=F.relu(self.deconv2(x))
x=self.deconv1(x)
return torch.sigmoid(x)
def forward(self, x):
mu, logvar, s = self.encode(x)
z = self.reparameterize(mu, logvar)
return self.decode(z, s), mu, logvar
def loss_function(recon_x, x, mu, logvar):
MSE = F.mse_loss(recon_x, x, reduction="sum")
BCE = F.binary_cross_entropy(recon_x, x, reduction='sum')
KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
return BCE +KLD, BCE, KLD, MSE
def train(train_dl, model, epoch, inp_im, out_im, lr1):
torch.backends.cudnn.benchmark = True
optimizer=torch.optim.Adam(model.parameters(), lr=lr1)
model.train()
train_loss=0.0
bce_loss=0.0
kld_loss=0.0
mse_loss=0.0
for idx, (data, label) in enumerate(train_dl):
data, label= Variable(data), Variable(label)
data1=data[:,0,:,:]
data1 = data1.unsqueeze(1)
data1 = data1.to(device)
out, mu, var=model(data1)
loss,bce,kld,mse=loss_function(out,data1, mu, var)
train_loss+=loss.item()
bce_loss+=bce.item()
kld_loss+=kld.item()
mse_loss+=mse.item()
model.zero_grad()
loss.backward()
optimizer.step()
av_loss= train_loss / len(train_dl.dataset)
av_kld_loss=kld_loss / len(train_dl.dataset)
av_bce_loss=bce_loss / len(train_dl.dataset)
av_mse_loss=mse_loss / len(train_dl.dataset)
return av_loss, av_bce_loss, av_kld_loss, av_mse_loss
def test(epoch, model):
model.eval()
test_loss = 0
test_mse_loss = 0
with torch.no_grad():
for i, (data, _) in enumerate(test_loader):
data = data.to(device)
data1=data[:,0,:,:]
data1 = data1.unsqueeze(1)
data1 = data1.to(device)
recon_batch, mu, logvar = model(data1)
loss,bce,kld,mse= loss_function(recon_batch, data1, mu, logvar)
test_loss +=loss.item()
test_mse_loss +=mse.item()
test_loss /= len(test_loader.dataset)
test_mse_loss /= len(test_loader.dataset)
return test_loss, test_mse_loss
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='VAE MNIST Example')
parser.add_argument('--batch-size', type=int, default=128, metavar='N',
help='input batch size for training (default: 128)')
parser.add_argument('-ks','--kernel-size', type=int, default=5, metavar='N',
help='kernel size (convolution)')
parser.add_argument('--epochs', type=int, default=200, metavar='N',
help='number of epochs to train (default: 10)')
parser.add_argument('-lr','--learning-rate', type=float, default=1e-3, metavar='',
help='learning rate (default: 1e-3)')
args = parser.parse_args()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
kwargs = {'num_workers': 1, 'pin_memory': True} if torch.cuda.is_available() else {}
transform1=transforms.Compose([transforms.ToTensor()])
train1=ImageFolder("/home/zaianir/Documents/dataset_s2/train", transform1)
test1=ImageFolder("/home/zaianir/Documents/dataset_s2/test", transform1)
train_loader=torch.utils.data.DataLoader(train1, batch_size=args.batch_size, shuffle=True, **kwargs)
test_loader=torch.utils.data.DataLoader(test1, batch_size=args.batch_size, shuffle=False, **kwargs)
data , target = next(iter(train_loader))
inp_s=data.shape
model = VAE(inp_s, args.kernel_size).to(device)
for epoch in range(1, args.epochs + 1):
moy_train_loss,av_bce_loss, av_kld_loss, av_mse_loss =train(train_loader, model, epoch, args.input_img, args.output_img, args.learning_rate)
test_l, mse_test_l=test(epoch, model)
I’m getting this error:
load conda/4.5.4 : OK
Traceback (most recent call last):
File “/home/uz/zaianir/scratch/VAE22.py”, line 194, in
moy_train_loss,av_bce_loss, av_kld_loss, av_mse_loss =train(train_loader, model, epoch, args.input_img, args.output_img, args.learning_rate)
File “/home/uz/zaianir/scratch/VAE22.py”, line 96, in train
loss.backward()
File “/home/uz/zaianir/.conda/envs/myenv2/lib/python3.7/site-packages/torch/tensor.py”, line 102, in backward
torch.autograd.backward(self, gradient, retain_graph, create_graph)
File “/home/uz/zaianir/.conda/envs/myenv2/lib/python3.7/site-packages/torch/autograd/init.py”, line 90, in backward
allow_unreachable=True) # allow_unreachable flag
RuntimeError: CUDA out of memory. Tried to allocate 7.64 GiB (GPU 0; 31.72 GiB total capacity; 20.76 GiB already allocated; 5.01 GiB free; 3.61 GiB cached)
- I’m running my validation code with
torch.no_grad() - I added
torch.backends.cudnn.benchmark = Truebefore training - I decreased the batch size from 128 to 20
But I still get the same error.
Can you help me please.
Thank you.