Model accuracy significantly decreases after applying of random transformation to the input data

Hello! I’m trying to imitate actions of bot on the game field. Input data is tensor 27x32x32 and action label (move to one of four directions or don’t move: 0, 1, 2, 3 or 4). I want to randomly transform data by rotating input data or flipping it horizontally or vertically with appropriate action change. But when I do that, accuracy of my model drops and doesn’t exceed 55% after 5 epochs. Without transformation it’s usually 75-77% after 5 epochs. What am I doing wrong?

Here is my code.

# dict for appropriate change of moving action during transformation
# 0 - no transform, 1 - rotation on 90, 2 - rotation on 180, 
# 3 - rotation on 270, 4 - horizontal flip, 5 - vertical flip
transform_dict = {0: {0:0, 1:1, 2:2, 3:3, 4:4},
                  1: {0:3, 1:2, 2:0, 3:1, 4:4}, # N -> E, S -> W, W -> N, E -> S
                  2: {0:1, 1:0, 2:3, 3:2, 4:4}, # N -> S, S -> N, W -> E, E -> W
                  3: {0:2, 1:3, 2:1, 3:0, 4:4}, # N -> W, S -> E, W -> S, E -> N
                  4: {0:0, 1:1, 2:3, 3:2, 4:4}, # N -> N, S -> S, W -> E, E -> W
                  5: {0:1, 1:0, 2:2, 3:3, 4:4}}  # N -> S, S -> N, W -> W, E -> E

class RandomChoice(torch.nn.Module):
    def __init__(self, transforms):
       super().__init__()
       self.transforms = transforms

    def __call__(self, x):
        idx = random.choice([i for i in range(len(self.transforms))])
        t = self.transforms[idx]
        x = torch.from_numpy(x)
        return t(x), idx

transform=RandomChoice([lambda x:x, 
                        lambda x:torch.rot90(x, 1, [2, 1]),
                        lambda x:torch.rot90(x, 2, [2, 1]),
                        lambda x:torch.rot90(x, 1, [1, 2]),
                        RandomHorizontalFlip(1), 
                        RandomVerticalFlip(1)])

# Dataset with random data transformation
class LDataset(Dataset):
    def __init__(self, obses, samples, transform=transform, transform_dict=transform_dict):
        self.obses = obses
        self.samples = samples
        self.data_len = len(self.samples)
        self.len = self.data_len
        self.transform = transform
        self.transform_dict = transform_dict
            
    def __len__(self):
        return self.len

    def __getitem__(self, idx):
        data_idx = idx % self.data_len
        # get data and convert it to 27x32x32 tensor
        obs_id, unit_id, action = self.samples[data_idx]
        obs = self.obses[obs_id]
        state = make_input(obs, unit_id)
        # transform data
        if self.transform:
            t = self.transform(state)
            state, action = t[0], self.transform_dict[t[1]][action]
        return state, action
    
# CNN
class BasicConv2d(nn.Module):
    def __init__(self, input_dim, output_dim, kernel_size, bn):
        super().__init__()
        self.conv = nn.Conv2d(
            input_dim, output_dim, 
            kernel_size=kernel_size, 
            padding=(kernel_size[0] // 2, kernel_size[1] // 2)
        )
        self.bn = nn.BatchNorm2d(output_dim) if bn else None

    def forward(self, x):
        h = self.conv(x)
        h = self.bn(h) if self.bn is not None else h
        return h


class LNet(nn.Module):
    def __init__(self):
        super().__init__()
        layers, filters = 12, 40
        self.conv0 = BasicConv2d(27, filters, (3, 3), True)
        self.blocks = nn.ModuleList([BasicConv2d(filters, filters, (3, 3), True) for _ in range(layers)])
        self.head_p = nn.Linear(filters, 5, bias=False)

    def forward(self, x):
        h = F.relu_(self.conv0(x))
        for block in self.blocks:
            h = F.relu_(h + block(h))
        h_head = (h * x[:,:1]).view(h.size(0), h.size(1), -1).sum(-1)
        p = self.head_p(h_head)
        return p

# function for model training
def train_model(model, dataloaders_dict, criterion, optimizer, num_epochs, city=False):
    
    best_acc = 0.0
    
    for epoch in range(num_epochs):
        model.cuda()
        
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()
            else:
                model.eval()
                
            epoch_loss = 0.0
            epoch_acc = 0
            
            dataloader = dataloaders_dict[phase]
            for item in tqdm(dataloader, leave=False):
                
                states = item[0].cuda().float()
                actions = item[1].cuda().long()

                optimizer.zero_grad()
                
                with torch.set_grad_enabled(phase == 'train'):
                    policy = model(states)
                    loss = criterion(policy, actions)
                    _, preds = torch.max(policy, 1)

                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                    epoch_loss += loss.item() * len(policy)
                    epoch_acc += torch.sum(preds == actions.data)

            data_size = len(dataloader.dataset)
            epoch_loss = epoch_loss / data_size
            epoch_acc = epoch_acc.double() / data_size
            
            print(f'Epoch {epoch + 1}/{num_epochs} | {phase:^5} | Loss: {epoch_loss:.4f} | Acc: {epoch_acc:.4f}')
        
        if epoch_acc > best_acc:
            traced = torch.jit.trace(model.cpu(), torch.rand(1, 27, 32, 32))
            traced.save('model.pth')
            best_acc = epoch_acc
            
# make train and val data loader
train, val = train_test_split(samples, test_size=0.1, random_state=42, stratify=labels)
batch_size = 128

train_loader = DataLoader(
    LuxDataset(obses, train), 
    batch_size=batch_size, 
    shuffle=True, 
    num_workers=2
)
val_loader = DataLoader(
    LuxDataset(obses, val), 
    batch_size=batch_size, 
    shuffle=False, 
    num_workers=2
)

# set NN parameters
model = LNet()
dataloaders_dict = {"train": train_loader, "val": val_loader}
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
num_epochs = 5

# train model
train_model(model, dataloaders_dict, criterion, optimizer, num_epochs=num_epochs)

Training results before transformations

Epoch 1/5 | train | Loss: 0.8268 | Acc: 0.6641
Epoch 1/5 |  val  | Loss: 0.7259 | Acc: 0.7068
Epoch 2/5 | train | Loss: 0.6621 | Acc: 0.7342
Epoch 2/5 |  val  | Loss: 0.6507 | Acc: 0.7401
Epoch 3/5 | train | Loss: 0.6066 | Acc: 0.7574
Epoch 3/5 |  val  | Loss: 0.6352 | Acc: 0.7480
Epoch 4/5 | train | Loss: 0.5706 | Acc: 0.7721
Epoch 4/5 |  val  | Loss: 0.5877 | Acc: 0.7673
Epoch 5/5 | train | Loss: 0.5446 | Acc: 0.7830
Epoch 5/5 |  val  | Loss: 0.5939 | Acc: 0.7645

After

Epoch 1/5 | train | Loss: 1.1865 | Acc: 0.4607
Epoch 1/5 |  val  | Loss: 1.1271 | Acc: 0.4939
Epoch 2/5 | train | Loss: 1.0952 | Acc: 0.5081
Epoch 2/5 |  val  | Loss: 1.0836 | Acc: 0.5109
Epoch 3/5 | train | Loss: 1.0632 | Acc: 0.5231
Epoch 3/5 |  val  | Loss: 1.0636 | Acc: 0.5243
Epoch 4/5 | train | Loss: 1.0433 | Acc: 0.5321
Epoch 4/5 |  val  | Loss: 1.0370 | Acc: 0.5364
Epoch 5/5 | train | Loss: 1.0268 | Acc: 0.5386
Epoch 5/5 |  val  | Loss: 1.0245 | Acc: 0.5403